From 2a01f2e8a62d85b4cdc8771eedda8799676c0a90 Mon Sep 17 00:00:00 2001
From: zephyr <i@zephyr.moe>
Date: Sun, 22 May 2022 09:24:45 +0900
Subject: [PATCH] allow custom alpn

---
 Cargo.lock          |  2 +-
 kaminari/Cargo.toml |  2 +-
 kaminari/src/mix.rs |  1 +
 kaminari/src/opt.rs | 25 ++++++++++++++++++-------
 kaminari/src/tls.rs | 24 ++++++++++++++++++++++--
 5 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 99f5149..7be3b95 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -129,7 +129,7 @@ dependencies = [
 
 [[package]]
 name = "kaminari"
-version = "0.8.0"
+version = "0.8.3"
 dependencies = [
  "lazy_static",
  "lightws",
diff --git a/kaminari/Cargo.toml b/kaminari/Cargo.toml
index c9d0dba..27442d2 100644
--- a/kaminari/Cargo.toml
+++ b/kaminari/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "kaminari"
-version = "0.8.2"
+version = "0.8.3"
 edition = "2021"
 authors = ["zephyr <i@zephyr.moe>"]
 keywords = ["lightws", "network"]
diff --git a/kaminari/src/mix.rs b/kaminari/src/mix.rs
index 6a4e124..8ab5521 100644
--- a/kaminari/src/mix.rs
+++ b/kaminari/src/mix.rs
@@ -288,6 +288,7 @@ mod test {
             }),
             tls: Some(TlsClientConf {
                 sni: String::from("abc"),
+                alpn: vec![Vec::from("h2"), Vec::from("http/1.1")],
                 insecure: true,
                 early_data: true,
             }),
diff --git a/kaminari/src/opt.rs b/kaminari/src/opt.rs
index 5ac782c..29b5968 100644
--- a/kaminari/src/opt.rs
+++ b/kaminari/src/opt.rs
@@ -58,12 +58,21 @@ pub fn get_tls_client_conf(s: &str) -> Option<TlsClientConf> {
     }
 
     let sni = get_opt!(it.clone(), "sni");
+    let alpn = get_opt!(it.clone(), "alpn");
     let insecure = has_opt!(it.clone(), "insecure");
     let early_data = has_opt!(it.clone(), "0rtt");
 
     if let Some(sni) = sni {
+        let alpn = alpn.map_or(Vec::new(), |s| {
+            s.split(',')
+                .map(str::trim)
+                .map(Vec::from)
+                .filter(|v| !v.is_empty())
+                .collect()
+        });
         Some(TlsClientConf {
             sni: String::from(sni),
+            alpn,
             insecure,
             early_data,
         })
@@ -148,10 +157,12 @@ mod test {
     #[test]
     fn tls_client_conf() {
         macro_rules! y {
-            ( $( ($s:expr, $sni: expr, $insecure: expr, $early_data: expr); )+ )=> {
+            ( $( ($s:expr, $sni: expr, $alpn: expr, $insecure: expr, $early_data: expr); )+ )=> {
                 $(
                     assert_eq!(get_tls_client_conf($s), Some(TlsClientConf{
                         sni: String::from($sni),
+                        alpn: $alpn.split(',').map(str::trim).map(Vec::from)
+                        .filter(|v|!v.is_empty()).collect(),
                         insecure: $insecure,
                         early_data: $early_data,
                     }));
@@ -160,12 +171,12 @@ mod test {
         }
 
         y![
-            ("tls;sni=a.b.c", "a.b.c", false, false);
-            ("tls;sni=a.b.c;", "a.b.c", false, false);
-            ("tls;sni=a.b.c;insecure", "a.b.c", true, false);
-            ("tls;sni=a.b.c;insecure;", "a.b.c", true, false);
-            ("tls;sni=a.b.c;insecure;0rtt", "a.b.c", true, true);
-            ("tls;sni=a.b.c;insecure;0rtt;", "a.b.c", true, true);
+            ("tls;sni=a.b.c", "a.b.c", "", false, false);
+            ("tls;sni=a.b.c;alpn=h2", "a.b.c", "h2", false, false);
+            ("tls;sni=a.b.c;alpn=http/1.1;insecure", "a.b.c", "http/1.1", true, false);
+            ("tls;sni=a.b.c;alpn=h2,http/1.1;insecure;", "a.b.c", "h2,http/1.1", true, false);
+            ("tls;sni=a.b.c;alpn=h2,http/1.1;insecure;0rtt", "a.b.c", "h2,http/1.1", true, true);
+            ("tls;sni=a.b.c;alpn=h2,http/1.1;insecure;0rtt;" ,"a.b.c", "h2,http/1.1", true, true);
         ];
     }
 
diff --git a/kaminari/src/tls.rs b/kaminari/src/tls.rs
index 0c21186..1b78836 100644
--- a/kaminari/src/tls.rs
+++ b/kaminari/src/tls.rs
@@ -18,16 +18,32 @@ pub use tokio_rustls::server::TlsStream as TlsServerStream;
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct TlsClientConf {
     pub sni: String,
+    pub alpn: Vec<Vec<u8>>,
     pub insecure: bool,
     pub early_data: bool,
 }
 
 impl Display for TlsClientConf {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        let mut alpn = [0u8; 64];
+        let mut cursor = 1;
+        alpn[0] = b'[';
+        for (i, b) in self.alpn.iter().enumerate() {
+            alpn[cursor..cursor + b.len()].copy_from_slice(b);
+            cursor += b.len();
+            if i != self.alpn.len() - 1 {
+                alpn[cursor..cursor + 2].copy_from_slice(b", ");
+                cursor += 2;
+            }
+        }
+        alpn[cursor] = b']';
+
+        let alpn = std::str::from_utf8(&alpn[..cursor + 1]).unwrap();
+
         write!(
             f,
-            "sni: {}, insecure: {}, early_data: {}",
-            self.sni, self.insecure, self.early_data
+            "sni: {}, alpn: {}, insecure: {}, early_data: {}",
+            self.sni, alpn, self.insecure, self.early_data
         )
     }
 }
@@ -62,6 +78,7 @@ impl<T> TlsConnect<T> {
     pub fn new(conn: T, conf: TlsClientConf) -> Self {
         let TlsClientConf {
             sni,
+            alpn,
             insecure,
             early_data,
         } = conf;
@@ -80,6 +97,7 @@ impl<T> TlsConnect<T> {
         };
 
         conf.enable_early_data = early_data;
+        conf.alpn_protocols = alpn;
 
         Self {
             conn,
@@ -92,6 +110,7 @@ impl<T> TlsConnect<T> {
     pub fn new_shared(conn: T, conf: TlsClientConf) -> Self {
         let TlsClientConf {
             sni,
+            alpn,
             insecure,
             early_data,
         } = conf;
@@ -104,6 +123,7 @@ impl<T> TlsConnect<T> {
             .with_no_client_auth();
 
         conf.enable_early_data = early_data;
+        conf.alpn_protocols = alpn;
 
         Self {
             conn,