add test for mixed

Author: vaporif

programs/solana/programs/ics26-router/src/instructions/ack_packet.rs

@@ -144,14 +144,15 @@ pub fn ack_packet<'info>(
 
     let payload = packet::get_single_payload(&packet)?;
 
-    let total_payload_chunks = total_payload_chunks(&msg.packet.payloads);
+    let total_payload_chunks = total_payload_chunks(&msg.packet.payloads)?;
     let proof_data = chunking::assemble_proof_chunks(chunking::AssembleProofParams {
         remaining_accounts: ctx.remaining_accounts,
         relayer: &ctx.accounts.relayer,
         submitter: ctx.accounts.relayer.key(),
         client_id: &msg.packet.source_client,
         sequence: msg.packet.sequence,
         delivery: &msg.proof.data,
+        // proof chunks come after payload chunks
         start_index: total_payload_chunks,
     })?;
 

programs/solana/programs/ics26-router/src/instructions/recv_packet.rs

@@ -162,7 +162,7 @@ pub fn recv_packet<'info>(
 
     let payload = packet::get_single_payload(&packet)?;
 
-    let total_payload_chunks = total_payload_chunks(&msg.packet.payloads);
+    let total_payload_chunks = total_payload_chunks(&msg.packet.payloads)?;
 
     let proof_data = chunking::assemble_proof_chunks(chunking::AssembleProofParams {
         remaining_accounts: ctx.remaining_accounts,
@@ -171,6 +171,7 @@ pub fn recv_packet<'info>(
         client_id: &msg.packet.dest_client,
         sequence: msg.packet.sequence,
         delivery: &msg.proof.data,
+        // proof chunks come after payload chunks
         start_index: total_payload_chunks,
     })?;
 

programs/solana/programs/ics26-router/src/instructions/timeout_packet.rs

@@ -143,14 +143,15 @@ pub fn timeout_packet<'info>(
 
     let payload = packet::get_single_payload(&packet)?;
 
-    let total_payload_chunks = total_payload_chunks(&msg.packet.payloads);
+    let total_payload_chunks = total_payload_chunks(&msg.packet.payloads)?;
     let proof_data = chunking::assemble_proof_chunks(chunking::AssembleProofParams {
         remaining_accounts: ctx.remaining_accounts,
         relayer: &ctx.accounts.relayer,
         submitter: ctx.accounts.relayer.key(),
         client_id: &msg.packet.source_client,
         sequence: msg.packet.sequence,
         delivery: &msg.proof.data,
+        // proof chunks come after payload chunks
         start_index: total_payload_chunks,
     })?;
 

programs/solana/programs/ics26-router/src/utils/chunking.rs

@@ -151,14 +151,22 @@ pub fn assemble_single_payload_chunks(params: AssemblePayloadParams) -> Result<V
     Ok(payload_data)
 }
 
-pub fn total_payload_chunks(payloads: &[MsgPayload]) -> usize {
-    payloads
+pub fn total_payload_chunks(payloads: &[MsgPayload]) -> Result<usize> {
+    require!(
+        payloads.windows(2).all(|w| {
+            matches!(
+                (&w[0].data, &w[1].data),
+                (Delivery::Inline { .. }, Delivery::Inline { .. })
+                    | (Delivery::Chunked { .. }, Delivery::Chunked { .. })
+            )
+        }),
+        RouterError::MixedDeliveryModes
+    );
+
+    Ok(payloads
         .iter()
-        .filter_map(|p| match &p.data {
-            Delivery::Chunked { total_chunks } => Some(*total_chunks as usize),
-            Delivery::Inline { .. } => None,
-        })
-        .sum()
+        .map(|p| p.data.total_chunks() as usize)
+        .sum())
 }
 
 /// Filter out chunk accounts from `remaining_accounts` before passing to IBC app CPI
@@ -416,6 +424,68 @@ pub fn validate_and_reconstruct_packet(
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_total_payload_chunks_rejects_mixed_delivery() {
+        let payloads = vec![
+            MsgPayload {
+                source_port: "port".to_string(),
+                dest_port: "port".to_string(),
+                version: "1".to_string(),
+                encoding: "json".to_string(),
+                data: Delivery::Inline {
+                    data: b"data".to_vec(),
+                },
+            },
+            MsgPayload {
+                source_port: "port".to_string(),
+                dest_port: "port".to_string(),
+                version: "1".to_string(),
+                encoding: "json".to_string(),
+                data: Delivery::Chunked { total_chunks: 2 },
+            },
+        ];
+
+        let result = total_payload_chunks(&payloads);
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_total_payload_chunks_all_inline() {
+        let payloads = vec![MsgPayload {
+            source_port: "port".to_string(),
+            dest_port: "port".to_string(),
+            version: "1".to_string(),
+            encoding: "json".to_string(),
+            data: Delivery::Inline {
+                data: b"data".to_vec(),
+            },
+        }];
+
+        assert_eq!(total_payload_chunks(&payloads).unwrap(), 0);
+    }
+
+    #[test]
+    fn test_total_payload_chunks_all_chunked() {
+        let payloads = vec![
+            MsgPayload {
+                source_port: "port".to_string(),
+                dest_port: "port".to_string(),
+                version: "1".to_string(),
+                encoding: "json".to_string(),
+                data: Delivery::Chunked { total_chunks: 2 },
+            },
+            MsgPayload {
+                source_port: "port".to_string(),
+                dest_port: "port".to_string(),
+                version: "1".to_string(),
+                encoding: "json".to_string(),
+                data: Delivery::Chunked { total_chunks: 3 },
+            },
+        ];
+
+        assert_eq!(total_payload_chunks(&payloads).unwrap(), 5);
+    }
+
     #[test]
     fn test_validate_and_reconstruct_packet_inline_mode_success() {
         let packet = solana_ibc_types::MsgPacket {