chore: Use real origin format in tests and assert X-Cll-Origin header is sent

Author: andrejrakic

rust/crates/sdk/src/stream.rs

@@ -400,36 +400,20 @@ mod tests {
 
     #[test]
     fn test_parse_origins_from_header_with_braces() {
-        let result = parse_origins_from_header(
-            "{wss://ws1.dataengine.chain.link,wss://ws2.dataengine.chain.link}",
-        );
-        assert_eq!(
-            result,
-            vec![
-                "wss://ws1.dataengine.chain.link".to_string(),
-                "wss://ws2.dataengine.chain.link".to_string(),
-            ]
-        );
+        let result = parse_origins_from_header("{001,002}");
+        assert_eq!(result, vec!["001".to_string(), "002".to_string()]);
     }
 
     #[test]
     fn test_parse_origins_from_header_without_braces() {
-        let result = parse_origins_from_header(
-            "wss://ws1.dataengine.chain.link,wss://ws2.dataengine.chain.link",
-        );
-        assert_eq!(
-            result,
-            vec![
-                "wss://ws1.dataengine.chain.link".to_string(),
-                "wss://ws2.dataengine.chain.link".to_string(),
-            ]
-        );
+        let result = parse_origins_from_header("001,002");
+        assert_eq!(result, vec!["001".to_string(), "002".to_string()]);
     }
 
     #[test]
     fn test_parse_origins_from_header_single_origin() {
-        let result = parse_origins_from_header("{wss://ws1.dataengine.chain.link}");
-        assert_eq!(result, vec!["wss://ws1.dataengine.chain.link".to_string()]);
+        let result = parse_origins_from_header("{001}");
+        assert_eq!(result, vec!["001".to_string()]);
     }
 
     #[test]

rust/crates/sdk/tests/stream_integration_tests.rs

@@ -263,6 +263,59 @@ async fn test_stream_ha_reconnect_merge() {
     assert_eq!(stats.deduplicated, expected_deduplicated);
 }
 
+#[tokio::test]
+async fn test_stream_ha_x_cll_origin_header() {
+    let mock_server = MockWebSocketServer::new("127.0.0.1:0").await;
+    let ws_url = format!("ws://{}", mock_server.address());
+
+    // Use distinct opaque origin IDs matching the real protocol format (e.g. {001,002}).
+    mock_server
+        .set_ha_origins(vec!["001".to_string(), "002".to_string()])
+        .await;
+
+    let config = Config::new(
+        "mock_key".to_string(),
+        "mock_secret".to_string(),
+        "mock_rest_url".to_string(),
+        ws_url,
+    )
+    .with_ws_ha(WebSocketHighAvailability::Enabled)
+    .with_ws_max_reconnect(MAX_RECONNECT_ATTEMPTS)
+    .build()
+    .expect("Failed to build config");
+
+    let mut stream = Stream::new(&config, vec![])
+        .await
+        .expect("Failed to create stream");
+
+    stream.listen().await.expect("Failed to start listening");
+
+    // Allow time for both WebSocket connections to be established.
+    sleep(Duration::from_millis(500)).await;
+
+    let received = mock_server.get_received_cll_origins().await;
+
+    // Assert 1: X-Cll-Origin header is present on every HA WebSocket connection.
+    assert_eq!(received.len(), 2, "Expected 2 WebSocket connections in HA mode");
+    for origin in &received {
+        assert!(
+            origin.is_some(),
+            "X-Cll-Origin header was missing on a WebSocket connection"
+        );
+    }
+
+    // Assert 2: The server sees distinct origin values matching the configured origins.
+    let mut actual: Vec<String> = received.into_iter().flatten().collect();
+    actual.sort();
+    assert_eq!(
+        actual,
+        vec!["001".to_string(), "002".to_string()],
+        "X-Cll-Origin header values did not match the configured origins"
+    );
+
+    stream.close().await.expect("Failed to close stream");
+}
+
 #[tokio::test]
 #[ignore] // Ignored because it takes a while to complete. To run it, use this command: cargo test -- --ignored
 async fn test_stream_ha_max_reconnection_attempts() {

rust/crates/sdk/tests/utils/mock_websocket_server.rs

@@ -3,9 +3,12 @@ use std::sync::Arc;
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
     net::{TcpListener, TcpStream},
-    sync::{mpsc, Mutex, Notify},
+    sync::{mpsc, oneshot, Mutex, Notify},
+};
+use tokio_tungstenite::{
+    accept_hdr_async,
+    tungstenite::{handshake::server::Request as WsRequest, protocol::Message},
 };
-use tokio_tungstenite::{accept_async, tungstenite::protocol::Message};
 
 enum ServerCommand {
     Send(Vec<u8>),
@@ -20,6 +23,9 @@ pub struct MockWebSocketServer {
     /// Origins returned in X-Cll-Available-Origins HEAD response.
     /// When None, defaults to two copies of the server's own ws:// address.
     ha_origins: Arc<Mutex<Option<Vec<String>>>>,
+    /// X-Cll-Origin header values captured from incoming WebSocket upgrade requests.
+    /// Some(value) if the header was present, None if absent.
+    received_cll_origins: Arc<Mutex<Vec<Option<String>>>>,
 }
 
 impl MockWebSocketServer {
@@ -35,10 +41,13 @@ impl MockWebSocketServer {
         let clients = Arc::new(Mutex::new(Vec::new()));
         let shutdown_notify = Arc::new(Notify::new());
         let ha_origins: Arc<Mutex<Option<Vec<String>>>> = Arc::new(Mutex::new(None));
+        let received_cll_origins: Arc<Mutex<Vec<Option<String>>>> =
+            Arc::new(Mutex::new(Vec::new()));
 
         let clients_accept = clients.clone();
         let shutdown_accept = shutdown_notify.clone();
         let ha_origins_accept = ha_origins.clone();
+        let received_accept = received_cll_origins.clone();
         let server_address = address.clone();
 
         tokio::spawn(async move {
@@ -55,7 +64,8 @@ impl MockWebSocketServer {
                                     ])
                                 };
                                 let clients_clone = clients_accept.clone();
-                                tokio::spawn(handle_connection(stream, origins, clients_clone));
+                                let received_clone = received_accept.clone();
+                                tokio::spawn(handle_connection(stream, origins, clients_clone, received_clone));
                             }
                             Err(e) => {
                                 println!("Error accepting connection: {:?}", e);
@@ -95,6 +105,7 @@ impl MockWebSocketServer {
             command_sender,
             shutdown_notify,
             ha_origins,
+            received_cll_origins,
         }
     }
 
@@ -122,15 +133,22 @@ impl MockWebSocketServer {
     pub async fn set_ha_origins(&self, origins: Vec<String>) {
         *self.ha_origins.lock().await = Some(origins);
     }
+
+    /// Returns the X-Cll-Origin header values captured from all WebSocket upgrade requests.
+    /// Some(value) means the header was present; None means it was absent.
+    pub async fn get_received_cll_origins(&self) -> Vec<Option<String>> {
+        self.received_cll_origins.lock().await.clone()
+    }
 }
 
 async fn handle_connection(
     mut stream: TcpStream,
     ha_origins: Vec<String>,
     clients: Arc<Mutex<Vec<mpsc::Sender<Message>>>>,
+    received_cll_origins: Arc<Mutex<Vec<Option<String>>>>,
 ) {
     // Peek at first 4 bytes to distinguish HTTP HEAD from WebSocket upgrade.
-    // peek() does not consume data, so the full request remains readable by accept_async.
+    // peek() does not consume data, so the full request remains readable by accept_hdr_async.
     let mut peek_buf = [0u8; 4];
     let n = match stream.peek(&mut peek_buf).await {
         Ok(n) => n,
@@ -152,14 +170,31 @@ async fn handle_connection(
         );
         let _ = stream.write_all(response.as_bytes()).await;
     } else {
-        // WebSocket upgrade
-        let ws_stream = match accept_async(stream).await {
+        // WebSocket upgrade — capture the X-Cll-Origin header from the upgrade request.
+        let (origin_tx, mut origin_rx) = oneshot::channel::<Option<String>>();
+
+        let ws_stream = match accept_hdr_async(stream, move |req: &WsRequest, resp| {
+            let origin = req
+                .headers()
+                .get("x-cll-origin")
+                .and_then(|v| v.to_str().ok())
+                .map(|s| s.to_string());
+            let _ = origin_tx.send(origin);
+            Ok(resp)
+        })
+        .await
+        {
             Ok(s) => s,
             Err(e) => {
                 println!("WebSocket accept error: {:?}", e);
                 return;
             }
         };
+
+        // origin_tx.send() has already run by the time accept_hdr_async resolves.
+        let origin = origin_rx.try_recv().unwrap_or(None);
+        received_cll_origins.lock().await.push(origin);
+
         println!(
             "Client connected: {}",
             ws_stream.get_ref().peer_addr().unwrap()