SSE streaming with long-poll fallback; encrypted unicast profile; registry enc keys

src/node.rs

@@ -38,6 +38,8 @@ pub struct Node {
     members: RwLock<Vec<Member>>,
     members_mtime: Mutex<Option<SystemTime>>,
     registry: Option<Arc<Watcher>>,
+    enc_secret: Option<String>,
+    enc_public: Option<String>,
     aggregates: Mutex<Aggregates>,
     seen: Mutex<HashSet<String>>,
     pending: Mutex<HashMap<String, Vec<(String, ResponseBody)>>>,
@@ -136,6 +138,11 @@ impl Node {
         if let Some(watcher) = &registry {
             watcher.load_initial();
         }
+        let enc_secret = config.load_enc_key()?;
+        let enc_public = match &enc_secret {
+            Some(secret) => Some(crate::crypto::enc_public_from_secret(secret)?),
+            None => None,
+        };
         Ok(Arc::new(Self {
             config,
             key,
@@ -143,6 +150,8 @@ impl Node {
             members: RwLock::new(members),
             members_mtime: Mutex::new(members_mtime),
             registry,
+            enc_secret,
+            enc_public,
             aggregates: Mutex::new(Aggregates::default()),
             seen: Mutex::new(HashSet::new()),
             pending: Mutex::new(HashMap::new()),
@@ -386,7 +395,52 @@ impl Node {
         delivered
     }
 
-    async fn dispatch(self: &Arc<Self>, envelope: Envelope, relay: &str) {
+    fn encrypt_for(&self, recipient_key: &str, body: &Value) -> Value {
+        let Some(watcher) = &self.registry else {
+            return body.clone();
+        };
+        let Some(recipient_enc) = watcher.enc_key(recipient_key) else {
+            return body.clone();
+        };
+        let context = crate::crypto::unicast_context(&self.identifier(), &recipient_enc);
+        match serde_json::to_vec(body) {
+            Ok(plaintext) => crate::crypto::encrypt_unicast(&recipient_enc, &context, &plaintext)
+                .unwrap_or_else(|_| body.clone()),
+            Err(_) => body.clone(),
+        }
+    }
+
+    fn decrypt_body(&self, envelope: &mut Envelope) -> Result<()> {
+        let Some(enc) = envelope.body.get("enc").cloned() else {
+            return Ok(());
+        };
+        let secret = self
+            .enc_secret
+            .as_deref()
+            .ok_or_else(|| anyhow!("encrypted message received but no enc key configured"))?;
+        let public = self
+            .enc_public
+            .as_deref()
+            .ok_or_else(|| anyhow!("encrypted message received but no enc key configured"))?;
+        let context = crate::crypto::unicast_context(&envelope.from, public);
+        let plaintext = crate::crypto::decrypt_unicast(secret, &context, &enc)?;
+        envelope.body = serde_json::from_slice(&plaintext).context("decrypted body is not json")?;
+        Ok(())
+    }
+
+    async fn send_payload(
+        &self,
+        to_key: &str,
+        msg_type: &str,
+        body: Value,
+        relay: &str,
+    ) -> Result<()> {
+        let body = self.encrypt_for(to_key, &body);
+        let envelope = Envelope::new(&self.key, &self.identifier(), msg_type, body);
+        self.send_unicast(to_key, envelope, relay).await
+    }
+
+    async fn dispatch(self: &Arc<Self>, mut envelope: Envelope, relay: &str) {
         if envelope.verify().is_err() {
             return;
         }
@@ -413,6 +467,10 @@ impl Node {
         if !listed {
             return;
         }
+        if let Err(error) = self.decrypt_body(&mut envelope) {
+            eprintln!("dropped encrypted message: {error}");
+            return;
+        }
         match envelope.msg_type.as_str() {
             TYPE_QUERY => {
                 if envelope.from == self.identifier() || !self.config.node.responder {
@@ -498,13 +556,8 @@ impl Node {
             return Ok(());
         }
         let body = build_response(&query.qid, response_items(&hits), total, max);
-        let envelope = Envelope::new(
-            &self.key,
-            &self.identifier(),
-            TYPE_RESPONSE,
-            serde_json::to_value(&body)?,
-        );
-        self.send_unicast(querier, envelope, relay).await
+        self.send_payload(querier, TYPE_RESPONSE, serde_json::to_value(&body)?, relay)
+            .await
     }
 
     pub async fn request_aggregate(
@@ -561,8 +614,10 @@ impl Node {
         let Ok(value) = serde_json::to_value(&body) else {
             return;
         };
-        let envelope = Envelope::new(&self.key, &self.identifier(), TYPE_AGGREGATE, value);
-        if let Err(error) = self.send_unicast(requester_key, envelope, relay).await {
+        if let Err(error) = self
+            .send_payload(requester_key, TYPE_AGGREGATE, value, relay)
+            .await
+        {
             eprintln!("aggregate reply failed: {error}");
         }
     }
@@ -618,6 +673,82 @@ async fn poll_relay(node: Arc<Node>, relay: String) {
             continue;
         };
         let signature = node.key.sign(&poll_signing_bytes(&member, &nonce));
+        let url = format!(
+            "{}/v1/stream?member={}&nonce={}&sig={}",
+            base, member, nonce, signature
+        );
+        match node.client.get(&url).send().await {
+            Ok(response) if response.status().is_success() => {
+                if let Err(error) = consume_stream(&node, &base, response).await {
+                    eprintln!("stream from {base} ended: {error}");
+                }
+                tokio::time::sleep(Duration::from_millis(200)).await;
+            }
+            Ok(response)
+                if response.status() == reqwest::StatusCode::NOT_FOUND
+                    || response.status() == reqwest::StatusCode::METHOD_NOT_ALLOWED =>
+            {
+                long_poll_relay(&node, &base).await;
+                return;
+            }
+            Ok(response) => {
+                eprintln!("stream from {base} rejected: {}", response.status());
+                tokio::time::sleep(Duration::from_secs(1)).await;
+            }
+            Err(_) => {
+                tokio::time::sleep(Duration::from_secs(1)).await;
+            }
+        }
+    }
+}
+
+async fn consume_stream(
+    node: &Arc<Node>,
+    base: &str,
+    response: reqwest::Response,
+) -> anyhow::Result<()> {
+    use futures_util::StreamExt;
+    let mut stream = response.bytes_stream();
+    let mut buffer = String::new();
+    let mut event = String::new();
+    let mut data = String::new();
+    while let Some(chunk) = stream.next().await {
+        let chunk = chunk?;
+        buffer.push_str(&String::from_utf8_lossy(&chunk));
+        while let Some(newline) = buffer.find('\n') {
+            let line = buffer[..newline].trim_end_matches('\r').to_string();
+            buffer.drain(..=newline);
+            if line.is_empty() {
+                if event == "envelope" && !data.is_empty() {
+                    if let Ok(envelope) = serde_json::from_str::<Envelope>(&data) {
+                        node.dispatch(envelope, base).await;
+                    }
+                } else if event == "lag" {
+                    eprintln!("relay {base} reports lag: {data}");
+                }
+                event.clear();
+                data.clear();
+            } else if let Some(rest) = line.strip_prefix("event:") {
+                event = rest.trim().to_string();
+            } else if let Some(rest) = line.strip_prefix("data:") {
+                if !data.is_empty() {
+                    data.push('\n');
+                }
+                data.push_str(rest.strip_prefix(' ').unwrap_or(rest));
+            }
+        }
+    }
+    Ok(())
+}
+
+async fn long_poll_relay(node: &Arc<Node>, base: &str) {
+    let member = node.key.public_hex();
+    loop {
+        let Some(nonce) = fetch_challenge(node, base, &member).await else {
+            tokio::time::sleep(Duration::from_secs(1)).await;
+            continue;
+        };
+        let signature = node.key.sign(&poll_signing_bytes(&member, &nonce));
         let url = format!(
             "{}/v1/poll?member={}&nonce={}&sig={}&timeout_ms=20000",
             base, member, nonce, signature
@@ -632,7 +763,7 @@ async fn poll_relay(node: Arc<Node>, relay: String) {
                                 if let Ok(envelope) =
                                     serde_json::from_value::<Envelope>(message.clone())
                                 {
-                                    node.dispatch(envelope, &base).await;
+                                    node.dispatch(envelope, base).await;
                                 }
                             }
                         }
@@ -696,6 +827,7 @@ async fn local_status(State(node): State<Arc<Node>>) -> Response {
     Json(json!({
         "name": node.config.node.name,
         "id": node.config.node.id,
+        "enc_key": node.enc_public,
         "registry_version": registry_version,
         "pubkey": node.key.public_hex(),
         "listen": node.config.node.listen,