From df69a54adf2e08b25b240a10debdb3ad63144463 Mon Sep 17 00:00:00 2001
From: Mike Dilger <mike@mikedilger.com>
Date: Mon, 15 Jul 2024 10:31:20 +1200
Subject: [PATCH] Honor Sec-WebSocket-Protocol:

- If it is missing, we speak nostr and we don't return it
- If it exists and contains "nostr", we speak nostr and we return it
- If it exists and doesn't contain "nostr" we throw a 501 NOT IMPLEMENTED StatusCode
---
 src/lib.rs | 32 +++++++++++++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/src/lib.rs b/src/lib.rs
index 2320212..654bdd7 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -155,6 +155,27 @@ async fn handle_http_request(
     }
 
     if hyper_tungstenite::is_upgrade_request(&request) {
+        // If the client asks for a Sec-Websocket-Protocol that we don't understand,
+        // Respond with 501 Not Implemented
+        let maybe_protocol: Option<String> = match request.headers().get("sec-websocket-protocol") {
+            None => None,
+            Some(hv) => hv.to_str().ok().map(|s| s.to_owned()),
+        };
+        if let Some(ref protocol) = maybe_protocol {
+            let mut we_can_do_nostr: bool = false;
+            for option in protocol.split(',') {
+                if option.trim() == "nostr" {
+                    we_can_do_nostr = true;
+                    break;
+                }
+            }
+            if !we_can_do_nostr {
+                return Ok(Response::builder()
+                    .status(StatusCode::NOT_IMPLEMENTED)
+                    .body(Full::new(Bytes::new()))?);
+            }
+        }
+
         let web_socket_config = WebSocketConfig {
             max_write_buffer_size: 1024 * 1024,  // 1 MB
             max_message_size: Some(1024 * 1024), // 1 MB
@@ -162,9 +183,18 @@ async fn handle_http_request(
             ..Default::default()
         };
 
-        let (response, websocket) =
+        let (mut response, websocket) =
             hyper_tungstenite::upgrade(&mut request, Some(web_socket_config))?;
 
+        // If the client asked for Sec-Websocket-Protocol, then we already checked it must
+        // have asked for 'nostr', so send that as a response header
+        if maybe_protocol.is_some() {
+            response.headers_mut().insert(
+                http::header::SEC_WEBSOCKET_PROTOCOL,
+                http::header::HeaderValue::from_static("nostr"),
+            );
+        }
+
         // Start the websocket thread
         tokio::spawn(async move { websocket_thread(peer, websocket, origin, ua).await });