Pass headers to websocket requests (#244)

Pass headers to websocket requests

---------

Co-authored-by: Alan Chen <[email protected]>

Author: benmanns

Cargo.lock

@@ -24,9 +24,9 @@ humantime-serde = { version = "1.1", default-features = false }
 tokio-stream = { version = "0.1", default-features = false, features = ["time"] }
 tokio = { version = "1.40", default-features = false, features = ["time"] }
 tokio-tungstenite = { version = "0.21.0", features = [ "rustls-tls-webpki-roots" ] }
-# tungstenite = { version = "0.21.0", features = ["rustls-tls-webpki-roots"] }
 encoding_rs = { version = "0.8", default-features = false }
 mime = { version = "0.3", default-features = false }
+http = "1.1"
 
 fluvio = { git = "https://github.com/infinyon/fluvio", branch = "wasmtime_21" }
 fluvio-connector-common = { git = "https://github.com/infinyon/fluvio", branch = "wasmtime_21", features = ["derive"] }

crates/http-source/Cargo.toml

@@ -14,7 +14,9 @@ use tokio::net::TcpStream;
 use tokio::time::Duration;
 use tokio_stream::{wrappers::IntervalStream, StreamExt};
 use tokio_tungstenite::{
-    connect_async, tungstenite::protocol::Message, MaybeTlsStream, WebSocketStream,
+    connect_async,
+    tungstenite::{client::IntoClientRequest, protocol::Message},
+    MaybeTlsStream, WebSocketStream,
 };
 use url::Url;
 
@@ -27,7 +29,7 @@ pub(crate) struct WebSocketSource {
 
 #[derive(Clone)]
 struct WSRequest {
-    url: Url,
+    request: tokio_tungstenite::tungstenite::handshake::client::Request,
     subscription_message: Option<String>,
 }
 
@@ -54,9 +56,9 @@ impl PingStream for WSPingOnlySink {
 }
 
 async fn establish_connection(request: WSRequest) -> Result<WebSocketStream<Transport>> {
-    match connect_async(&request.url).await {
+    match connect_async(request.request.clone()).await {
         Ok((mut ws_stream, _)) => {
-            info!("WebSocket connected to {}", &request.url);
+            info!("WebSocket connected to {}", &request.request.uri());
             if let Some(message) = request.subscription_message.as_ref() {
                 ws_stream.send(Message::Text(message.to_owned())).await?;
             }
@@ -128,10 +130,31 @@ async fn websocket_writer_and_stream<'a>(
 impl WebSocketSource {
     pub(crate) fn new(config: &HttpConfig) -> Result<Self> {
         let ws_config = config.websocket_config.as_ref();
+
+        let mut request = Url::parse(&config.endpoint.resolve()?)?.into_client_request()?;
+        let headers = request.headers_mut();
+
+        for h in config.headers.iter() {
+            match h.resolve() {
+                Ok(h) => {
+                    if let Some((key, value)) = h.split_once(':') {
+                        headers.insert(
+                            http::HeaderName::from_bytes(key.as_bytes())?,
+                            value.parse()?,
+                        );
+                    } else {
+                        error!("Failed to split header");
+                    }
+                }
+                Err(e) => {
+                    error!("Failed to resolve header: {}", e);
+                }
+            }
+        }
+
         Ok(Self {
             request: WSRequest {
-                url: Url::parse(&config.endpoint.resolve()?)
-                    .context("unable to parse http endpoint")?,
+                request,
                 subscription_message: ws_config.and_then(|c| c.subscription_message.to_owned()),
             },
             ping_interval_ms: ws_config.and_then(|c| c.ping_interval_ms).unwrap_or(10_000),

crates/http-source/src/websocket_source.rs

@@ -38,6 +38,21 @@ async fn main() -> tide::Result<()> {
             }
             Ok(())
         }));
+    app.at("/websocket-auth")
+        .get(WebSocket::new(|request, stream| async move {
+            let header_values = request.header("x-secret-token");
+            if header_values.is_none() || header_values.unwrap().last() != "abc123" {
+                stream.send_string("Unauthorized".to_string()).await?;
+                return Ok(());
+            }
+
+            for i in 1..11 {
+                stream
+                    .send_string(format!("Hello, Fluvio! - {}", i))
+                    .await?;
+            }
+            Ok(())
+        }));
 
     app.listen("127.0.0.1:8080").await?;
     Ok(())

crates/mock-http-server/src/main.rs

@@ -0,0 +1,12 @@
+meta:
+  version: latest
+  name: websocket-headers-connector
+  type: websocket-source
+  topic: TOPIC
+  create_topic: false
+  producer:
+    linger: 0ms
+http:
+  endpoint: ws://127.0.0.1:8080/websocket-auth
+  headers:
+    - "x-secret-token: abc123"

tests/websocket-headers-test-config.yaml

@@ -0,0 +1,41 @@
+#!/usr/bin/env bats
+
+setup() {
+    cargo build -p mock-http-server
+    ./target/debug/mock-http-server & disown
+    MOCK_PID=$!
+    FILE=$(mktemp)
+    cp ./tests/websocket-headers-test-config.yaml $FILE
+    UUID=$(uuidgen | awk '{print tolower($0)}')
+    TOPIC=${UUID}-topic
+    fluvio topic create $TOPIC
+
+    sed -i.BAK "s/TOPIC/${TOPIC}/g" $FILE
+    cat $FILE
+
+    cargo build -p http-source
+    ./target/debug/http-source --config $FILE & disown
+    CONNECTOR_PID=$!
+}
+
+teardown() {
+    fluvio topic delete $TOPIC
+    kill $MOCK_PID
+    kill $CONNECTOR_PID
+}
+
+@test "websocket-connector-test" {
+    count=1
+    echo "Starting consumer on topic $TOPIC"
+    sleep 13
+
+    fluvio consume -B -d $TOPIC | while read input; do
+        expected="Hello, Fluvio! - $count"
+        echo $input = $expected
+        [ "$input" = "$expected" ]
+        count=$(($count + 1))
+        if [ $count -eq 10 ]; then
+            break;
+        fi
+    done
+}