From 8dc8a413cbd91caf66be710fe357d76614f64d5a Mon Sep 17 00:00:00 2001
From: kazk <kazk.dev@gmail.com>
Date: Tue, 14 Sep 2021 19:50:27 -0700
Subject: [PATCH] Add extension negotiation methods to `WebSocketConfig`

---
 src/extensions/compression/deflate.rs | 11 ++++-----
 src/handshake/client.rs               |  5 ++--
 src/handshake/server.rs               |  4 ++--
 src/protocol/mod.rs                   | 34 +++++++++++++++++++++++++++
 4 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/src/extensions/compression/deflate.rs b/src/extensions/compression/deflate.rs
index 9f4930d..3310910 100644
--- a/src/extensions/compression/deflate.rs
+++ b/src/extensions/compression/deflate.rs
@@ -4,7 +4,7 @@ use flate2::{Compress, Compression, Decompress, FlushCompress, FlushDecompress,
 use http::HeaderValue;
 use thiserror::Error;
 
-use crate::{extensions, protocol::Role};
+use crate::protocol::Role;
 
 const PER_MESSAGE_DEFLATE: &str = "permessage-deflate";
 const CLIENT_NO_CONTEXT_TAKEOVER: &str = "client_no_context_takeover";
@@ -67,11 +67,10 @@ impl DeflateConfig {
         to_header_value(&offers)
     }
 
-    // This can be used for `WebSocket::from_raw_socket_with_compression`.
     /// Returns negotiation response based on offers and `DeflateContext` to manage per message compression.
-    pub fn negotiation_response<'a>(
+    pub(crate) fn accept_offer<'a>(
         &'a self,
-        extensions: impl Iterator<Item = &'a HeaderValue>,
+        offers: impl Iterator<Item = impl Iterator<Item = (&'a str, Option<&'a str>)>>,
     ) -> Option<(HeaderValue, DeflateContext)> {
         // Accept the first valid offer for `permessage-deflate`.
         // A server MUST decline an extension negotiation offer for this
@@ -83,9 +82,7 @@ impl DeflateConfig {
         // * The negotiation offer contains multiple extension parameters with
         //   the same name.
         // * The server doesn't support the offered configuration.
-        'outer: for (_, offer) in
-            extensions::iter_all(extensions).filter(|&(k, _)| k == self.name())
-        {
+        'outer: for offer in offers {
             let mut config =
                 DeflateConfig { compression: self.compression, ..DeflateConfig::default() };
             let mut agreed = Vec::new();
diff --git a/src/handshake/client.rs b/src/handshake/client.rs
index d428ca8..04c6a05 100644
--- a/src/handshake/client.rs
+++ b/src/handshake/client.rs
@@ -141,9 +141,8 @@ fn generate_request(
         }
         writeln!(req, "{}: {}\r", k, v.to_str()?).unwrap();
     }
-    if let Some(compression) = &config.and_then(|c| c.compression) {
-        let offer = compression.generate_offer();
-        writeln!(req, "Sec-WebSocket-Extensions: {}\r", offer.to_str()?).unwrap();
+    if let Some(offers) = config.and_then(|c| c.generate_offers()) {
+        writeln!(req, "Sec-WebSocket-Extensions: {}\r", offers.to_str()?).unwrap();
     }
     writeln!(req, "\r").unwrap();
     trace!("Request: {:?}", String::from_utf8_lossy(&req));
diff --git a/src/handshake/server.rs b/src/handshake/server.rs
index 4451441..7db9878 100644
--- a/src/handshake/server.rs
+++ b/src/handshake/server.rs
@@ -245,9 +245,9 @@ impl<S: Read + Write, C: Callback> HandshakeRole for ServerHandshake<S, C> {
                 }
 
                 let mut response = create_response(&result)?;
-                if let Some(compression) = &self.config.and_then(|c| c.compression) {
+                if let Some(config) = &self.config {
                     let extensions = result.headers().get_all("Sec-WebSocket-Extensions").iter();
-                    if let Some((agreed, pmce)) = compression.negotiation_response(extensions) {
+                    if let Some((agreed, pmce)) = config.accept_offers(extensions) {
                         self.pmce = Some(pmce);
                         response.headers_mut().insert("Sec-WebSocket-Extensions", agreed);
                     }
diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs
index b240a58..d998582 100644
--- a/src/protocol/mod.rs
+++ b/src/protocol/mod.rs
@@ -6,6 +6,7 @@ mod message;
 
 pub use self::{frame::CloseFrame, message::Message};
 
+use http::HeaderValue;
 use log::*;
 use std::{
     collections::VecDeque,
@@ -73,6 +74,39 @@ impl Default for WebSocketConfig {
     }
 }
 
+impl WebSocketConfig {
+    // Generate extension negotiation offers for configured extensions.
+    // Only `permessage-deflate` is supported at the moment.
+    pub(crate) fn generate_offers(&self) -> Option<HeaderValue> {
+        self.compression.map(|c| c.generate_offer())
+    }
+
+    // TODO Replace `DeflateContext` with something more general
+    // This can be used with `WebSocket::from_raw_socket_with_compression` for integration.
+    /// Returns negotiation response based on offers and `DeflateContext` to manage per message compression.
+    pub fn accept_offers<'a>(
+        &'a self,
+        extensions: impl Iterator<Item = &'a HeaderValue>,
+    ) -> Option<(HeaderValue, DeflateContext)> {
+        if let Some(compression) = &self.compression {
+            let extensions = crate::extensions::iter_all(extensions);
+            let offers =
+                extensions.filter_map(
+                    |(k, v)| {
+                        if k == compression.name() {
+                            Some(v)
+                        } else {
+                            None
+                        }
+                    },
+                );
+            compression.accept_offer(offers)
+        } else {
+            None
+        }
+    }
+}
+
 /// WebSocket input-output stream.
 ///
 /// This is THE structure you want to create to be able to speak the WebSocket protocol.