diff --git a/crates/rustack-sqs-core/src/message.rs b/crates/rustack-sqs-core/src/message.rs index 248c0b97..43edfdb7 100644 --- a/crates/rustack-sqs-core/src/message.rs +++ b/crates/rustack-sqs-core/src/message.rs @@ -28,7 +28,7 @@ pub struct QueueMessage { pub approximate_first_receive_timestamp: Option, /// FIFO-only: sequence number. pub sequence_number: Option, - /// FIFO-only: message group ID. + /// Message group ID: required on FIFO queues, optional on standard (fair) queues. pub message_group_id: Option, /// FIFO-only: message deduplication ID. pub message_deduplication_id: Option, diff --git a/crates/rustack-sqs-core/src/queue/actor.rs b/crates/rustack-sqs-core/src/queue/actor.rs index 4b1de8a7..ab8d338f 100644 --- a/crates/rustack-sqs-core/src/queue/actor.rs +++ b/crates/rustack-sqs-core/src/queue/actor.rs @@ -55,7 +55,13 @@ struct PendingTransfer { } impl RedriveManager { - fn handoff(&mut self, message: &QueueMessage, target_arn: &str, source_arn: &str) -> bool { + fn handoff( + &mut self, + message: &QueueMessage, + target_arn: &str, + source_arn: &str, + source_is_fifo: bool, + ) -> bool { if self.quiescing || self.pending.len() >= 128 || source_arn == target_arn { return false; } @@ -78,6 +84,7 @@ impl RedriveManager { .try_send(QueueCommand::Transfer { message: message.clone(), source_arn: source_arn.to_owned(), + source_is_fifo, reply, }) .is_err() @@ -127,6 +134,8 @@ pub enum QueueCommand { message: QueueMessage, /// Exact source queue ARN. source_arn: String, + /// Whether the source is a FIFO queue; a dead-letter queue must be the same type. + source_is_fifo: bool, /// Confirms successful enqueue or rejection before mutation. reply: oneshot::Sender>, }, @@ -411,12 +420,13 @@ impl QueueActor { QueueCommand::Transfer { message, source_arn, + source_is_fifo, reply, } => { if reply.is_closed() { return false; } - let result = self.accept_transfer(message, &source_arn); + let result = self.accept_transfer(message, &source_arn, source_is_fifo); let accepted = result.is_ok(); let _ = reply.send(result); accepted @@ -525,8 +535,9 @@ impl QueueActor { &mut self, mut message: QueueMessage, source_arn: &str, + source_is_fifo: bool, ) -> Result<(), SqsError> { - if self.stopping || self.is_fifo != message.message_group_id.is_some() { + if self.stopping || self.is_fifo != source_is_fifo { return Err(SqsError::invalid_parameter_value( "DLQ unavailable or queue type mismatch", )); @@ -650,6 +661,10 @@ impl QueueActor { )); } + if let Some(group_id) = &input.message_group_id { + validate_message_group_id(group_id)?; + } + if self.is_fifo { self.handle_send_message_fifo(input) } else { @@ -663,13 +678,8 @@ impl QueueActor { &mut self, input: SendMessageInput, ) -> Result<(SendMessageOutput, bool), SqsError> { - // Reject FIFO-only fields on standard queues. - if input.message_group_id.is_some() { - return Err(SqsError::invalid_parameter_value( - "Value for parameter MessageGroupId is invalid. Reason: The request includes a \ - parameter that is not valid for this queue type.", - )); - } + // Standard queues take a MessageGroupId since SQS fair queues. It steers delivery + // fairness, which this emulator does not model, but is still returned on receive. if input.message_deduplication_id.is_some() { return Err(SqsError::invalid_parameter_value( "Value for parameter MessageDeduplicationId is invalid. Reason: The request \ @@ -704,7 +714,7 @@ impl QueueActor { approximate_receive_count: 0, approximate_first_receive_timestamp: None, sequence_number: None, - message_group_id: None, + message_group_id: input.message_group_id, message_deduplication_id: None, available_at, delay_seconds, @@ -903,6 +913,7 @@ impl QueueActor { manager: &mut self.redrive, attributes: &self.attributes, source_arn: &self.arn, + source_is_fifo: self.is_fifo, }; match &mut self.storage { QueueStorage::Standard(storage) => try_receive_standard( @@ -1081,6 +1092,19 @@ struct RedriveRequest<'a> { manager: &'a mut RedriveManager, attributes: &'a QueueAttributes, source_arn: &'a str, + source_is_fifo: bool, +} + +/// SQS accepts 1 to 128 letters, digits and punctuation, that is printable ASCII without space, +/// as a MessageGroupId on either queue type. +fn validate_message_group_id(group_id: &str) -> Result<(), SqsError> { + if (1..=128).contains(&group_id.len()) && group_id.bytes().all(|byte| byte.is_ascii_graphic()) { + return Ok(()); + } + Err(SqsError::invalid_parameter_value(format!( + "Value {group_id} for parameter MessageGroupId is invalid. Reason: MessageGroupId can \ + only include alphanumeric and punctuation characters. 1 to 128 in length." + ))) } impl RedriveRequest<'_> { @@ -1116,7 +1140,12 @@ fn try_receive_standard( match storage.available.pop_front() { Some(mut msg) => { if let Some(target) = redrive.target(&msg) { - if !redrive.manager.handoff(&msg, &target, redrive.source_arn) { + if !redrive.manager.handoff( + &msg, + &target, + redrive.source_arn, + redrive.source_is_fifo, + ) { // Standard queues have no group ordering constraint: retain // the poison message without starving healthy messages. storage.available.push_back(msg); @@ -1166,7 +1195,10 @@ fn try_receive_fifo( for (mut msg, group_id) in received { if let Some(target) = redrive.target(&msg) { - if !redrive.manager.handoff(&msg, &target, redrive.source_arn) { + if !redrive + .manager + .handoff(&msg, &target, redrive.source_arn, redrive.source_is_fifo) + { storage.finish_redrive(&group_id, Some(msg)); } continue; @@ -1498,7 +1530,10 @@ fn merge_attribute_names(old: &[String], new: &[String]) -> Vec { mod tests { use std::collections::HashMap; - use rustack_sqs_model::input::{ReceiveMessageInput, SendMessageInput}; + use rustack_sqs_model::{ + error::SqsErrorCode, + input::{ReceiveMessageInput, SendMessageInput}, + }; use tokio::sync::{mpsc, oneshot}; use super::{QueueActor, QueueCommand}; @@ -1583,6 +1618,100 @@ mod tests { ); } + #[test] + fn test_should_keep_message_group_id_on_standard_queue_for_receive() { + let mut actor = test_actor(); + let (reply, mut response) = oneshot::channel(); + + actor.handle_command(QueueCommand::SendMessage { + input: SendMessageInput { + message_group_id: Some("tenant-a".to_owned()), + ..send_input(None) + }, + reply, + }); + assert!( + response + .try_recv() + .expect("send command should reply") + .is_ok() + ); + + let (reply, mut response) = oneshot::channel(); + actor.handle_command(QueueCommand::ReceiveMessage { + input: ReceiveMessageInput { + message_system_attribute_names: vec!["MessageGroupId".to_owned()], + ..receive_input(0) + }, + reply, + }); + + let output = response + .try_recv() + .expect("receive command should reply") + .expect("receive should succeed"); + assert_eq!(output.messages.len(), 1); + assert_eq!( + output.messages[0] + .attributes + .get("MessageGroupId") + .map(String::as_str), + Some("tenant-a") + ); + } + + #[test] + fn test_should_check_message_group_id_format_on_standard_and_fifo_queues() { + let (_sender, receiver) = mpsc::channel(8); + let fifo = QueueActor::new( + "test-queue.fifo".to_owned(), + "arn:aws:sqs:us-east-1:000000000000:test-queue.fifo".to_owned(), + true, + QueueAttributes { + content_based_deduplication: true, + ..QueueAttributes::default() + }, + receiver, + HashMap::new(), + "000000000000".to_owned(), + 0, + ); + let longest = "a".repeat(128); + let too_long = "a".repeat(129); + let cases = [ + ("tenant-a", true), + ("!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~", true), + (longest.as_str(), true), + ("", false), + (too_long.as_str(), false), + ("has space", false), + ("caf\u{e9}", false), + ("tab\t", false), + ]; + + for mut actor in [test_actor(), fifo] { + for (group_id, valid) in cases { + let (reply, mut response) = oneshot::channel(); + actor.handle_command(QueueCommand::SendMessage { + input: SendMessageInput { + message_group_id: Some(group_id.to_owned()), + message_body: format!("{group_id} body"), + ..send_input(None) + }, + reply, + }); + let result = response.try_recv().expect("send command should reply"); + match result { + Ok(_) => assert!(valid, "accepted {group_id:?} on fifo={}", actor.is_fifo), + Err(error) => { + assert!(!valid, "rejected {group_id:?}: {error:?}"); + assert_eq!(error.code, SqsErrorCode::InvalidParameterValue); + } + } + } + } + } + #[test] fn immediate_send_fulfills_existing_long_poll() { let mut actor = test_actor(); diff --git a/crates/rustack-sqs-core/src/redrive_tests.rs b/crates/rustack-sqs-core/src/redrive_tests.rs index 93649f65..e44e1745 100644 --- a/crates/rustack-sqs-core/src/redrive_tests.rs +++ b/crates/rustack-sqs-core/src/redrive_tests.rs @@ -81,6 +81,36 @@ async fn test_should_deliver_standard_and_fifo_dead_letters_to_real_target() { } } +#[tokio::test] +async fn test_should_deliver_standard_dead_letter_with_message_group_id() { + let provider = RustackSqs::new(SqsConfig::default()); + let target = create(&provider, "fair-dead", None).await; + let source = create(&provider, "fair-source", Some("fair-dead")).await; + provider + .send_message(SendMessageInput { + queue_url: source.clone(), + message_body: "grouped".into(), + message_group_id: Some("tenant-a".into()), + ..Default::default() + }) + .await + .unwrap(); + assert_eq!(receive(&provider, &source).await.len(), 1); + assert!(receive(&provider, &source).await.is_empty()); + provider.quiesce().await.unwrap(); + let messages = receive(&provider, &target).await; + assert_eq!(messages.len(), 1); + assert_eq!( + messages[0] + .attributes + .get("MessageGroupId") + .map(String::as_str), + Some("tenant-a") + ); + assert!(receive(&provider, &source).await.is_empty()); + provider.shutdown_all().await; +} + #[tokio::test] async fn test_should_retain_dead_letters_until_missing_target_becomes_available() { for suffix in ["", ".fifo"] { diff --git a/tests/integration/src/test_sqs.rs b/tests/integration/src/test_sqs.rs index 21f2af75..29b6130c 100644 --- a/tests/integration/src/test_sqs.rs +++ b/tests/integration/src/test_sqs.rs @@ -834,16 +834,15 @@ mod tests { #[tokio::test] #[ignore = "requires running server"] - async fn test_should_reject_fifo_fields_on_standard_queue() { + async fn test_should_reject_deduplication_id_on_standard_queue() { let client = sqs_client(); let (_, url) = create_test_queue(&client, "no-fifo").await; - // Sending with MessageGroupId on a standard queue should fail. let err = client .send_message() .queue_url(&url) .message_body("test") - .message_group_id("group-a") + .message_deduplication_id("dedup-a") .send() .await; @@ -852,6 +851,44 @@ mod tests { delete_test_queue(&client, &url).await; } + #[tokio::test] + #[ignore = "requires running server"] + async fn test_should_accept_message_group_id_on_standard_queue() { + let client = sqs_client(); + let (_, url) = create_test_queue(&client, "fair-queue").await; + + // Fair queues: a standard queue takes a MessageGroupId and returns it on receive. + client + .send_message() + .queue_url(&url) + .message_body("test") + .message_group_id("tenant-a") + .send() + .await + .expect("send with MessageGroupId to a standard queue"); + + let received = client + .receive_message() + .queue_url(&url) + .message_system_attribute_names( + aws_sdk_sqs::types::MessageSystemAttributeName::MessageGroupId, + ) + .send() + .await + .expect("receive"); + let message = &received.messages()[0]; + assert_eq!( + message + .attributes() + .and_then(|attributes| attributes + .get(&aws_sdk_sqs::types::MessageSystemAttributeName::MessageGroupId)) + .map(String::as_str), + Some("tenant-a"), + ); + + delete_test_queue(&client, &url).await; + } + // --------------------------------------------------------------------------- // Bug: Long-poll does not wake up when a message is sent mid-poll // --------------------------------------------------------------------------- diff --git a/tests/sqs-compat/test_errors.py b/tests/sqs-compat/test_errors.py index 0de3f26e..d67686e2 100644 --- a/tests/sqs-compat/test_errors.py +++ b/tests/sqs-compat/test_errors.py @@ -78,16 +78,32 @@ def test_should_reject_message_too_large(self, sqs, test_queue): class TestFifoErrors: - def test_should_reject_group_id_on_standard_queue(self, sqs, test_queue): - """MessageGroupId should not be accepted on standard queues.""" - with pytest.raises(botocore.exceptions.ClientError) as exc: - sqs.send_message( - QueueUrl=test_queue, - MessageBody="test", - MessageGroupId="group", - ) - err = exc.value.response["Error"]["Code"] - assert err in ("InvalidParameterValue", "InvalidAddress") + def test_should_accept_group_id_on_standard_queue(self, sqs, test_queue): + """Standard (fair) queues accept MessageGroupId and return it on receive.""" + sqs.send_message( + QueueUrl=test_queue, + MessageBody="test", + MessageGroupId="group", + ) + resp = sqs.receive_message( + QueueUrl=test_queue, + MessageSystemAttributeNames=["MessageGroupId"], + WaitTimeSeconds=1, + ) + messages = resp.get("Messages", []) + assert len(messages) == 1 + assert messages[0]["Attributes"]["MessageGroupId"] == "group" + + def test_should_reject_invalid_group_id_on_standard_queue(self, sqs, test_queue): + """MessageGroupId takes 1 to 128 letters, digits and punctuation.""" + for group_id in ("has space", "x" * 129): + with pytest.raises(botocore.exceptions.ClientError) as exc: + sqs.send_message( + QueueUrl=test_queue, + MessageBody="test", + MessageGroupId=group_id, + ) + assert exc.value.response["Error"]["Code"] == "InvalidParameterValue" def test_should_reject_dedup_id_on_standard_queue(self, sqs, test_queue): """MessageDeduplicationId should not be accepted on standard queues."""