fix(server-ng): add messages encryption (#3644)

This PR adds support for `messages` encryption in `server-ng` binary.

Author: numinnex

core/integration/tests/server/scenarios/encryption_scenario.rs

@@ -93,6 +93,10 @@ async fn should_fill_data_with_headers_and_verify_after_restart_using_api(encryp
         .await
         .unwrap();
 
+    // server-ng has no flush primitive (FLUSH_UNSAVED_BUFFER denies typed);
+    // the eager-flush envs in `build_server_config` make every committed
+    // batch hit disk instead.
+    #[cfg(not(feature = "vsr"))]
     client
         .flush_unsaved_buffer(
             &Identifier::named(stream_name).unwrap(),
@@ -103,7 +107,7 @@ async fn should_fill_data_with_headers_and_verify_after_restart_using_api(encryp
         .await
         .unwrap();
 
-    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+    tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
 
     // Verify on-disk encryption of headers and payload
     let data_path = harness.server().data_path();
@@ -247,6 +251,10 @@ async fn should_fill_data_with_headers_and_verify_after_restart_using_api(encryp
         .await
         .unwrap();
 
+    // server-ng has no flush primitive (FLUSH_UNSAVED_BUFFER denies typed);
+    // the eager-flush envs in `build_server_config` make every committed
+    // batch hit disk instead.
+    #[cfg(not(feature = "vsr"))]
     client
         .flush_unsaved_buffer(
             &Identifier::named(stream_name).unwrap(),
@@ -257,7 +265,7 @@ async fn should_fill_data_with_headers_and_verify_after_restart_using_api(encryp
         .await
         .unwrap();
 
-    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+    tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
 
     let polled = client
         .poll_messages(
@@ -490,6 +498,17 @@ fn encryption_disabled() -> bool {
 fn build_server_config(encryption: bool) -> TestServerConfig {
     let mut extra_envs = HashMap::new();
 
+    // server-ng flushes on the journal thresholds (no flush primitive), so
+    // force every committed batch straight to disk for the on-disk asserts.
+    extra_envs.insert(
+        "IGGY_SYSTEM_PARTITION_MESSAGES_REQUIRED_TO_SAVE".to_string(),
+        "1".to_string(),
+    );
+    extra_envs.insert(
+        "IGGY_SYSTEM_PARTITION_ENFORCE_FSYNC".to_string(),
+        "true".to_string(),
+    );
+
     if encryption {
         extra_envs.insert(
             "IGGY_SYSTEM_ENCRYPTION_ENABLED".to_string(),

core/integration/tests/server/scenarios/mod.rs

@@ -34,7 +34,6 @@ pub mod create_message_payload;
 // shard-0 HTTP listener and the create/delete commit through the metadata STM,
 // so the token replicates to every shard a TCP client may land on.
 pub mod cross_protocol_pat_scenario;
-#[cfg(not(feature = "vsr"))]
 pub mod encryption_scenario;
 pub mod invalid_consumer_offset_scenario;
 // Asserts server log-file rotation/archival policies; server-ng's file

core/partitions/src/iggy_partitions.rs

@@ -23,9 +23,10 @@ use crate::{IggyPartition, Partition, PollingArgs, PollingConsumer};
 use ahash::AHashSet;
 use consensus::{Consensus, Plane, PlaneIdentity, VsrConsensus};
 use iggy_binary_protocol::{
-    Command2, ConsensusHeader, PrepareHeader, PrepareOkHeader, RequestHeader,
+    Command2, ConsensusHeader, Operation, PrepareHeader, PrepareOkHeader, RequestHeader,
 };
 use message_bus::MessageBus;
+use server_common::send_messages2::{convert_request_message, encrypt_batch_request};
 use server_common::sharding::{IggyNamespace, LocalIdx, ShardId};
 #[cfg(debug_assertions)]
 use std::cell::Cell;
@@ -468,6 +469,32 @@ where
             );
             return;
         }
+        // At-rest encryption happens HERE, once, before the op enters
+        // consensus: canonicalize the wire form first so both the legacy and
+        // v2 request encodings encrypt identically, then encrypt payload +
+        // user headers per message. The ciphertext is what gets journaled,
+        // replicated, checksummed, and persisted -- every replica stores
+        // identical bytes -- and the poll reply is the single decrypt point.
+        let message = if message.header().operation == Operation::SendMessages
+            && let Some(encryptor) = &self.config().encryptor
+        {
+            let canonical = convert_request_message(namespace, message)
+                .and_then(|message| encrypt_batch_request(message, encryptor));
+            match canonical {
+                Ok(message) => message,
+                Err(error) => {
+                    warn!(
+                        target: "iggy.partitions.diag",
+                        namespace_raw = namespace.inner(),
+                        %error,
+                        "dropping send_messages: failed to encrypt batch at ingestion"
+                    );
+                    return;
+                }
+            }
+        } else {
+            message
+        };
         let Some(partition) = self.get_mut_by_ns(&namespace) else {
             warn!(
                 target: "iggy.partitions.diag",

core/partitions/src/types.rs

@@ -15,9 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use iggy_common::{IggyByteSize, PollingStrategy};
+use iggy_common::{EncryptorKind, IggyByteSize, PollingStrategy};
 use server_common::iobuf::Frozen;
 use smallvec::SmallVec;
+use std::sync::Arc;
 
 #[derive(Debug, Clone)]
 pub struct Fragment<const ALIGN: usize = 4096> {
@@ -216,6 +217,13 @@ pub struct PartitionsConfig {
     pub enforce_fsync: bool,
     /// Maximum size of a single segment before rotation.
     pub segment_size: IggyByteSize,
+    /// Server-side at-rest encryption. Applied ONCE, on the primary at
+    /// ingestion, so the ciphertext replicates verbatim: every replica
+    /// journals, acks, and persists identical bytes (checksums and the
+    /// deterministic segment rolls both depend on that), and the poll path
+    /// decrypts uniformly whether a fragment came from the resident journal
+    /// or from disk.
+    pub encryptor: Option<Arc<EncryptorKind>>,
 }
 
 impl PartitionsConfig {

core/server-ng/src/bootstrap.rs

@@ -46,7 +46,7 @@ use iggy_common::defaults::{
     DEFAULT_ROOT_USERNAME, MAX_PASSWORD_LENGTH, MAX_USERNAME_LENGTH, MIN_PASSWORD_LENGTH,
     MIN_USERNAME_LENGTH,
 };
-use iggy_common::{IggyByteSize, PartitionStats, variadic};
+use iggy_common::{Aes256GcmEncryptor, EncryptorKind, IggyByteSize, PartitionStats, variadic};
 use journal::Journal;
 use journal::prepare_journal::PrepareJournal;
 use message_bus::client_listener::{self, RequestHandler};
@@ -1414,6 +1414,16 @@ async fn build_shard_for_thread(
     let owned_partitions_capacity = total_partitions
         .div_ceil(usize::from(total_shards).max(1))
         .saturating_mul(2);
+    // At-rest encryption: built once per shard from the shared config; the
+    // ingestion path encrypts on the primary and the poll reply decrypts.
+    // A bad key fails the boot rather than silently serving plaintext.
+    let encryptor = if config.system.encryption.enabled {
+        let aes = Aes256GcmEncryptor::from_base64_key(&config.system.encryption.key)
+            .map_err(|error| ServerNgError::Iggy(Box::new(error)))?;
+        Some(Arc::new(EncryptorKind::Aes256Gcm(aes)))
+    } else {
+        None
+    };
     let partitions = IggyPartitions::with_capacity(
         shard_local_id,
         PartitionsConfig {
@@ -1424,6 +1434,7 @@ async fn build_shard_for_thread(
                 .size_of_messages_required_to_save,
             enforce_fsync: config.system.partition.enforce_fsync,
             segment_size: config.system.segment.size,
+            encryptor,
         },
         owned_partitions_capacity,
     );

core/server-ng/src/dispatch.rs

@@ -1512,15 +1512,20 @@ async fn handle_poll_messages(
                 Some(PartitionReadReply::Poll {
                     fragments,
                     current_offset,
-                }) => build_polled_messages_body(partition_id, current_offset, fragments)
-                    .unwrap_or_else(|error| {
-                        warn!(
-                            transport_client_id,
-                            error = %error,
-                            "failed to re-encode polled batches; replying empty poll"
-                        );
-                        empty_polled_messages_body(partition_id)
-                    }),
+                }) => build_polled_messages_body(
+                    partition_id,
+                    current_offset,
+                    fragments,
+                    shard.plane.partitions().config().encryptor.as_deref(),
+                )
+                .unwrap_or_else(|error| {
+                    warn!(
+                        transport_client_id,
+                        error = %error,
+                        "failed to re-encode polled batches; replying empty poll"
+                    );
+                    empty_polled_messages_body(partition_id)
+                }),
                 other => {
                     warn!(
                         transport_client_id,

core/server-ng/src/http/handlers.rs

@@ -27,7 +27,7 @@ use axum::extract::{Path, Query, State};
 use axum::http::{HeaderMap, HeaderName, HeaderValue, StatusCode, header};
 use axum::response::{IntoResponse, Response};
 use chrono::Local;
-use consensus::MetadataHandle;
+use consensus::{MetadataHandle, PartitionsHandle};
 use iggy_binary_protocol::codes::{
     GET_CONSUMER_GROUP_CODE, GET_CONSUMER_GROUPS_CODE, GET_PERSONAL_ACCESS_TOKENS_CODE,
     GET_STATS_CODE, GET_STREAM_CODE, GET_STREAMS_CODE, GET_TOPIC_CODE, GET_TOPICS_CODE,
@@ -1033,8 +1033,13 @@ pub(in crate::http) async fn poll_messages(
             fragments,
             current_offset,
         }) => {
-            let body = build_polled_messages_body(partition_id, current_offset, fragments)
-                .map_err(ReadError::Rejected)?;
+            let body = build_polled_messages_body(
+                partition_id,
+                current_offset,
+                fragments,
+                state.shard.plane.partitions().config().encryptor.as_deref(),
+            )
+            .map_err(ReadError::Rejected)?;
             Ok(Json(
                 PolledMessages::from_bytes(body).map_err(ReadError::Rejected)?,
             ))

core/server-ng/src/http/wire.rs

@@ -509,9 +509,13 @@ mod tests {
         header.base_offset = 41;
         header.base_timestamp = 999_999;
 
-        let body =
-            build_polled_messages_body(3, 42, fragment_from_stored_batch(&header, &stored.blob))
-                .expect("re-encodes wire body");
+        let body = build_polled_messages_body(
+            3,
+            42,
+            fragment_from_stored_batch(&header, &stored.blob),
+            None,
+        )
+        .expect("re-encodes wire body");
         let polled = PolledMessages::from_bytes(body).expect("decodes as the SDK does");
 
         assert_eq!(polled.partition_id, 3);

core/server-ng/src/partition_reconciler.rs

@@ -1054,6 +1054,7 @@ mod tests {
                 size_of_messages_required_to_save: iggy_common::IggyByteSize::from(1024_u64),
                 enforce_fsync: false,
                 segment_size: config.system.segment.size,
+                encryptor: None,
             },
         );
         let shards_table = PapayaShardsTable::new();

core/server-ng/src/responses.rs

@@ -77,7 +77,7 @@ use iggy_binary_protocol::{
     Command2, GenericHeader, IGGY_PROTOCOL_VERSION, KIND_CONSUMER_GROUP, Operation, ReplyHeader,
     RequestHeader, WireDecode, WireEncode, WireIdentifier, WireName, WirePartitioning,
 };
-use iggy_common::{Identifier, IggyError, IggyTimestamp, MaxTopicSize};
+use iggy_common::{EncryptorKind, Identifier, IggyError, IggyTimestamp, MaxTopicSize};
 use metadata::impls::metadata::StreamsFrontend;
 use partitions::PollFragments;
 use server_common::Message;
@@ -1307,6 +1307,7 @@ pub(crate) fn build_polled_messages_body(
     partition_id: u32,
     current_offset: u64,
     fragments: PollFragments,
+    encryptor: Option<&EncryptorKind>,
 ) -> Result<Bytes, IggyError> {
     // Body head: [partition_id:4][current_offset:8][count:4]. `count` sits at
     // COUNT_OFFSET and is backpatched once the walk below knows it.
@@ -1374,20 +1375,56 @@ pub(crate) fn build_polled_messages_body(
             body.extend_from_slice(&offset.to_le_bytes());
             body.extend_from_slice(&timestamp.to_le_bytes());
             body.extend_from_slice(&origin_timestamp.to_le_bytes());
-            body.extend_from_slice(
-                &u32::try_from(user_headers_length)
-                    .expect("length came from u32")
-                    .to_le_bytes(),
-            );
-            body.extend_from_slice(
-                &u32::try_from(payload_length)
-                    .expect("length came from u32")
-                    .to_le_bytes(),
-            );
-            body.extend_from_slice(&0u64.to_le_bytes()); // reserved
-            // Stored sections are already in legacy order
-            // (`[payload][user_headers]`): copy through contiguously.
-            body.extend_from_slice(&stream[sections_start..sections_end]);
+            if let Some(encryptor) = encryptor {
+                // At-rest encryption: stored sections are ciphertext (encrypted
+                // once at ingestion, replicated verbatim); this reply is the
+                // single decrypt point, so lengths are rewritten to the
+                // plaintext sizes. The stored per-message checksum still covers
+                // the ciphertext and is passed through untouched (the SDK does
+                // not re-validate it against the reply layout).
+                let payload_end = sections_start + payload_length;
+                let payload = encryptor
+                    .decrypt(&stream[sections_start..payload_end])
+                    .map_err(|_| IggyError::CannotDecryptData)?;
+                let user_headers = if user_headers_length > 0 {
+                    Some(
+                        encryptor
+                            .decrypt(&stream[payload_end..sections_end])
+                            .map_err(|_| IggyError::CannotDecryptData)?,
+                    )
+                } else {
+                    None
+                };
+                let user_headers_bytes: &[u8] = user_headers.as_deref().unwrap_or_default();
+                body.extend_from_slice(
+                    &u32::try_from(user_headers_bytes.len())
+                        .map_err(|_| IggyError::InvalidCommand)?
+                        .to_le_bytes(),
+                );
+                body.extend_from_slice(
+                    &u32::try_from(payload.len())
+                        .map_err(|_| IggyError::InvalidCommand)?
+                        .to_le_bytes(),
+                );
+                body.extend_from_slice(&0u64.to_le_bytes()); // reserved
+                body.extend_from_slice(&payload);
+                body.extend_from_slice(user_headers_bytes);
+            } else {
+                body.extend_from_slice(
+                    &u32::try_from(user_headers_length)
+                        .expect("length came from u32")
+                        .to_le_bytes(),
+                );
+                body.extend_from_slice(
+                    &u32::try_from(payload_length)
+                        .expect("length came from u32")
+                        .to_le_bytes(),
+                );
+                body.extend_from_slice(&0u64.to_le_bytes()); // reserved
+                // Stored sections are already in legacy order
+                // (`[payload][user_headers]`): copy through contiguously.
+                body.extend_from_slice(&stream[sections_start..sections_end]);
+            }
 
             count += 1;
             cursor = sections_end;