Skip to main content

ruststream_amqp/
message.rs

1//! [`AmqpMessage`] and the mapping between `RustStream` headers and `AMQP` 1.0 message sections.
2//!
3//! Well-known `RustStream` headers ride the `AMQP` `properties` section (`content-type`,
4//! `correlation-id`, `reply-to`, `message-id`, and the partition key as `group-id`); every other
5//! header rides `application-properties`, so no envelope format is invented and non-RustStream
6//! peers see plain `AMQP` messages.
7
8use 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
19/// Header carrying the partition key, mapped onto the `AMQP` `group-id` property.
20///
21/// Mirrors the in-memory broker's convention, so services can switch brokers without changing
22/// their headers.
23pub const PARTITION_KEY_HEADER: &str = "partition-key";
24
25/// How a delivered message asks its pump task to settle it.
26#[derive(Debug)]
27pub(crate) enum SettleKind {
28    /// Accept the delivery (ack).
29    Accept,
30    /// Release the delivery back to the broker for redelivery (nack with requeue).
31    Release,
32    /// Reject the delivery as undeliverable (nack without requeue); the broker's dead-letter
33    /// policy decides what happens next.
34    Reject,
35}
36
37/// A settlement request shipped from a message handle to the subscription's pump task.
38#[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
47/// A message delivered by an [`AmqpSubscriber`](crate::AmqpSubscriber).
48///
49/// `ack` maps to the `accept` disposition, `nack(requeue = true)` to `release`, and
50/// `nack(requeue = false)` to `reject` (terminal; the broker's dead-letter policy applies).
51/// Deliveries received on an at-most-once subscription are already settled, so `ack`/`nack`
52/// report [`AckError::Unsupported`] instead of pretending.
53pub struct AmqpMessage {
54    payload: Bytes,
55    headers: Headers,
56    /// `None` when the delivery is already settled (at-most-once, request/reply replies).
57    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
143/// Builds the `AMQP` message for an outgoing publish.
144pub(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
195/// Extracts `RustStream` headers from a delivered `AMQP` message.
196pub(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
223/// Renders any `AMQP` message-id form as text, so it survives the byte-valued header map.
224fn 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
242/// Scalar fallback: peers may put numbers or booleans into application-properties; text is the
243/// only lossless byte form the header map can carry for them.
244fn 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
264/// Extracts the payload bytes from a delivered body.
265///
266/// `Data` sections are the native byte path (what this crate publishes); a string or binary
267/// `AmqpValue` from a foreign peer is accepted as bytes too. Anything else has no faithful byte
268/// form and is reported as [`AmqpError::UnsupportedBody`].
269pub(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}