1use bytes::Bytes;
9use fe2o3_amqp::link::delivery::DeliveryInfo;
10use fe2o3_amqp_types::messaging::{
11 ApplicationProperties, Body, Data, Message, MessageId, Properties,
12};
13use fe2o3_amqp_types::primitives::{Binary, SimpleValue, Symbol, Value};
14use ruststream::{AckError, Headers, IncomingMessage, OutgoingMessage, Partitioned};
15use tokio::sync::{mpsc, oneshot};
16
17use crate::error::AmqpError;
18
19pub const PARTITION_KEY_HEADER: &str = "partition-key";
24
25#[derive(Debug)]
27pub(crate) enum SettleKind {
28 Accept,
30 Release,
32 Reject,
35}
36
37#[derive(Debug)]
39pub(crate) struct SettleCmd {
40 pub(crate) info: DeliveryInfo,
41 pub(crate) kind: SettleKind,
42 pub(crate) done: oneshot::Sender<Result<(), AckError>>,
43}
44
45pub(crate) type SettleSender = mpsc::UnboundedSender<SettleCmd>;
46
47pub struct AmqpMessage {
54 payload: Bytes,
55 headers: Headers,
56 settle: Option<SettleHandle>,
58}
59
60struct SettleHandle {
61 tx: SettleSender,
62 info: DeliveryInfo,
63}
64
65impl std::fmt::Debug for AmqpMessage {
66 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67 f.debug_struct("AmqpMessage")
68 .field("payload_len", &self.payload.len())
69 .field("settled", &self.settle.is_none())
70 .finish_non_exhaustive()
71 }
72}
73
74impl AmqpMessage {
75 pub(crate) fn unsettled(
76 payload: Bytes,
77 headers: Headers,
78 tx: SettleSender,
79 info: DeliveryInfo,
80 ) -> Self {
81 Self {
82 payload,
83 headers,
84 settle: Some(SettleHandle { tx, info }),
85 }
86 }
87
88 pub(crate) fn settled(payload: Bytes, headers: Headers) -> Self {
89 Self {
90 payload,
91 headers,
92 settle: None,
93 }
94 }
95
96 async fn settle(self, kind: SettleKind) -> Result<(), AckError> {
97 let Some(SettleHandle { tx, info }) = self.settle else {
98 return Err(AckError::Unsupported);
99 };
100 let (done, wait) = oneshot::channel();
101 tx.send(SettleCmd { info, kind, done }).map_err(|_| {
102 AckError::Broker(Box::from("the subscription's pump task has shut down"))
103 })?;
104 wait.await.map_err(|_| {
105 AckError::Broker(Box::from("the subscription's pump task has shut down"))
106 })?
107 }
108}
109
110impl Partitioned for AmqpMessage {
111 fn partition_key(&self) -> Option<&[u8]> {
112 self.headers.get(PARTITION_KEY_HEADER)
113 }
114}
115
116impl IncomingMessage for AmqpMessage {
117 fn payload(&self) -> &[u8] {
118 &self.payload
119 }
120
121 fn headers(&self) -> &Headers {
122 &self.headers
123 }
124
125 async fn ack(self) -> Result<(), AckError> {
126 self.settle(SettleKind::Accept).await
127 }
128
129 async fn nack(self, requeue: bool) -> Result<(), AckError> {
130 let kind = if requeue {
131 SettleKind::Release
132 } else {
133 SettleKind::Reject
134 };
135 self.settle(kind).await
136 }
137
138 fn partition_key(&self) -> Option<&[u8]> {
139 Partitioned::partition_key(self)
140 }
141}
142
143pub(crate) fn to_amqp_message(msg: &OutgoingMessage<'_>) -> Message<Data> {
145 let headers = msg.headers();
146 let mut properties = Properties::default();
147 let mut has_properties = false;
148 let mut application: Option<ApplicationProperties> = None;
149
150 for (name, value) in headers.iter() {
151 let text = || String::from_utf8_lossy(value).into_owned();
152 match name {
153 "content-type" => {
154 properties.content_type = Some(Symbol::from(text()));
155 has_properties = true;
156 }
157 "correlation-id" => {
158 properties.correlation_id = Some(MessageId::String(text()));
159 has_properties = true;
160 }
161 "reply-to" => {
162 properties.reply_to = Some(text());
163 has_properties = true;
164 }
165 "message-id" => {
166 properties.message_id = Some(MessageId::String(text()));
167 has_properties = true;
168 }
169 PARTITION_KEY_HEADER => {
170 properties.group_id = Some(text());
171 has_properties = true;
172 }
173 other => {
174 let simple = std::str::from_utf8(value).map_or_else(
175 |_| SimpleValue::Binary(Binary::from(value.to_vec())),
176 |s| SimpleValue::String(s.to_owned()),
177 );
178 application
179 .get_or_insert_with(ApplicationProperties::default)
180 .insert(other.to_owned(), simple);
181 }
182 }
183 }
184
185 let mut builder = Message::builder();
186 if has_properties {
187 builder = builder.properties(properties);
188 }
189 if let Some(application) = application {
190 builder = builder.application_properties(application);
191 }
192 builder.data(Binary::from(msg.payload().to_vec())).build()
193}
194
195pub(crate) fn headers_from_amqp<B>(message: &Message<B>) -> Headers {
197 let mut headers = Headers::new();
198 if let Some(properties) = &message.properties {
199 if let Some(content_type) = &properties.content_type {
200 headers.insert("content-type", content_type.to_string());
201 }
202 if let Some(correlation_id) = &properties.correlation_id {
203 headers.insert("correlation-id", message_id_text(correlation_id));
204 }
205 if let Some(reply_to) = &properties.reply_to {
206 headers.insert("reply-to", reply_to.clone());
207 }
208 if let Some(message_id) = &properties.message_id {
209 headers.insert("message-id", message_id_text(message_id));
210 }
211 if let Some(group_id) = &properties.group_id {
212 headers.insert(PARTITION_KEY_HEADER, group_id.clone());
213 }
214 }
215 if let Some(application) = &message.application_properties {
216 for (name, value) in application.iter() {
217 headers.insert(name.clone(), simple_value_bytes(value));
218 }
219 }
220 headers
221}
222
223fn message_id_text(id: &MessageId) -> String {
225 match id {
226 MessageId::String(s) => s.clone(),
227 MessageId::Uuid(u) => format!("{u:x}"),
228 MessageId::Ulong(n) => n.to_string(),
229 MessageId::Binary(b) => String::from_utf8_lossy(b).into_owned(),
230 }
231}
232
233fn simple_value_bytes(value: &SimpleValue) -> Bytes {
234 match value {
235 SimpleValue::String(s) => Bytes::copy_from_slice(s.as_bytes()),
236 SimpleValue::Binary(b) => Bytes::copy_from_slice(b),
237 SimpleValue::Symbol(s) => Bytes::copy_from_slice(s.as_str().as_bytes()),
238 other => Bytes::from(format_simple_value(other)),
239 }
240}
241
242fn format_simple_value(value: &SimpleValue) -> String {
245 match value {
246 SimpleValue::Bool(v) => v.to_string(),
247 SimpleValue::Ubyte(v) => v.to_string(),
248 SimpleValue::Ushort(v) => v.to_string(),
249 SimpleValue::Uint(v) => v.to_string(),
250 SimpleValue::Ulong(v) => v.to_string(),
251 SimpleValue::Byte(v) => v.to_string(),
252 SimpleValue::Short(v) => v.to_string(),
253 SimpleValue::Int(v) => v.to_string(),
254 SimpleValue::Long(v) => v.to_string(),
255 SimpleValue::Float(v) => v.to_string(),
256 SimpleValue::Double(v) => v.to_string(),
257 SimpleValue::Char(v) => v.to_string(),
258 SimpleValue::Timestamp(v) => v.milliseconds().to_string(),
259 SimpleValue::Uuid(v) => format!("{v:x}"),
260 other => format!("{other:?}"),
261 }
262}
263
264pub(crate) fn payload_from_body(body: Body<Value>, address: &str) -> Result<Bytes, AmqpError> {
270 match body {
271 Body::Data(batch) => {
272 let mut chunks = batch.into_iter();
273 match (chunks.next(), chunks.next()) {
274 (None, _) => Ok(Bytes::new()),
275 (Some(Data(first)), None) => Ok(Bytes::from(first.into_vec())),
276 (Some(Data(first)), Some(Data(second))) => {
277 let mut all = first.into_vec();
278 all.extend_from_slice(&second);
279 for Data(chunk) in chunks {
280 all.extend_from_slice(&chunk);
281 }
282 Ok(Bytes::from(all))
283 }
284 }
285 }
286 Body::Value(value) => match value.0 {
287 Value::Binary(b) => Ok(Bytes::from(b.into_vec())),
288 Value::String(s) => Ok(Bytes::from(s.into_bytes())),
289 _ => Err(AmqpError::UnsupportedBody {
290 address: address.to_owned(),
291 }),
292 },
293 Body::Empty => Ok(Bytes::new()),
294 Body::Sequence(_) => Err(AmqpError::UnsupportedBody {
295 address: address.to_owned(),
296 }),
297 }
298}
299
300#[cfg(test)]
301mod tests {
302 use super::*;
303
304 #[test]
305 fn well_known_headers_ride_the_properties_section() {
306 let mut headers = Headers::new();
307 headers.insert("content-type", "application/json");
308 headers.insert("correlation-id", "corr-1");
309 headers.insert("reply-to", "replies");
310 headers.insert("message-id", "msg-1");
311 headers.insert(PARTITION_KEY_HEADER, "user-42");
312 headers.insert("x-custom", "value");
313 let outgoing = OutgoingMessage::new("orders", b"{}".as_slice()).with_headers(headers);
314
315 let message = to_amqp_message(&outgoing);
316 let properties = message.properties.as_ref().expect("properties set");
317 assert_eq!(
318 properties.content_type.as_ref().map(Symbol::as_str),
319 Some("application/json")
320 );
321 assert_eq!(
322 properties.correlation_id,
323 Some(MessageId::String("corr-1".into()))
324 );
325 assert_eq!(properties.reply_to.as_deref(), Some("replies"));
326 assert_eq!(properties.group_id.as_deref(), Some("user-42"));
327 let application = message
328 .application_properties
329 .as_ref()
330 .expect("application properties set");
331 assert_eq!(
332 application.get("x-custom"),
333 Some(&SimpleValue::String("value".into()))
334 );
335 assert!(application.get("content-type").is_none());
336 }
337
338 #[test]
339 fn headers_round_trip_through_the_amqp_sections() {
340 let mut headers = Headers::new();
341 headers.insert("content-type", "application/json");
342 headers.insert("correlation-id", "corr-1");
343 headers.insert(PARTITION_KEY_HEADER, "user-42");
344 headers.insert("x-custom", "value");
345 let outgoing =
346 OutgoingMessage::new("orders", b"{}".as_slice()).with_headers(headers.clone());
347
348 let restored = headers_from_amqp(&to_amqp_message(&outgoing));
349 assert_eq!(restored.get_str("content-type"), Some("application/json"));
350 assert_eq!(restored.get_str("correlation-id"), Some("corr-1"));
351 assert_eq!(restored.get_str(PARTITION_KEY_HEADER), Some("user-42"));
352 assert_eq!(restored.get_str("x-custom"), Some("value"));
353 }
354
355 #[test]
356 fn data_body_yields_payload_bytes() {
357 let outgoing = OutgoingMessage::new("orders", b"payload".as_slice());
358 let message = to_amqp_message(&outgoing);
359 let body = Body::<Value>::Data(vec![message.body].into());
360 let payload = payload_from_body(body, "orders").expect("data body decodes");
361 assert_eq!(payload.as_ref(), b"payload");
362 }
363
364 #[test]
365 fn string_and_binary_values_are_accepted_as_bytes() {
366 let s = Body::Value(fe2o3_amqp_types::messaging::AmqpValue(Value::String(
367 "hi".into(),
368 )));
369 assert_eq!(
370 payload_from_body(s, "a").expect("string decodes").as_ref(),
371 b"hi"
372 );
373
374 let b = Body::Value(fe2o3_amqp_types::messaging::AmqpValue(Value::Binary(
375 Binary::from(b"raw".to_vec()),
376 )));
377 assert_eq!(
378 payload_from_body(b, "a").expect("binary decodes").as_ref(),
379 b"raw"
380 );
381 }
382
383 #[test]
384 fn foreign_value_bodies_are_reported_unsupported() {
385 let body = Body::Value(fe2o3_amqp_types::messaging::AmqpValue(Value::Bool(true)));
386 assert!(matches!(
387 payload_from_body(body, "a"),
388 Err(AmqpError::UnsupportedBody { .. })
389 ));
390 }
391}