feat: optional rustls support (#330)

* initial implementation of rustls support

* Refactor create_self_signed_cert.sh script

* resolve lint errors

* Fix handling of Option in tls.rs

* Update cargo-hack check command and feature dependencies

* fix missing point

* Add conditional check to skip test if client or server is not enabled

* clean up things

* fix for windows CI

* try fixing Windows CI

* Update src/main.rs

* Update src/transport/websocket.rs

* add missing messages

* split the tls mod

Co-authored-by: Ning Sun <n@sunng.info>

.github/workflows/rust.yml

@@ -32,7 +32,9 @@ jobs:
       - name: Setup cargo-hack
         run: cargo install cargo-hack
       - name: Check all features
-        run: cargo hack check --feature-powerset --no-dev-deps
+        run: >
+          cargo hack check --feature-powerset --no-dev-deps
+          --mutually-exclusive-features default,native-tls,websocket-native-tls,rustls,websocket-rustls
 
   build:
     name: Build for ${{ matrix.target }}
@@ -62,8 +64,10 @@ jobs:
       - uses: Swatinem/rust-cache@v1
       - name: Build
         run: cargo build
-      - name: Run tests
+      - name: Run tests with native-tls
         run: cargo test --verbose
+      - name: Run tests with rustls
+        run: cargo test --verbose --no-default-features --features server,client,rustls,noise,websocket-rustls,hot-reload
       - uses: actions/upload-artifact@v2
         with:
           name: rathole-${{ matrix.target }}

Cargo.lock

@@ -243,6 +243,15 @@ dependencies = [
  "generic-array",
 ]
 
+[[package]]
+name = "block-padding"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a8894febbff9f758034a5b8e12d87918f56dfc64a8e1fe757d65e29041538d93"
+dependencies = [
+ "generic-array",
+]
+
 [[package]]
 name = "byteorder"
 version = "1.4.3"
@@ -258,6 +267,15 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "cbc"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "26b52a9543ae338f279b96b0b9fed9c8093744685043739079ce85cd58f289a6"
+dependencies = [
+ "cipher",
+]
+
 [[package]]
 name = "cc"
 version = "1.0.83"
@@ -497,6 +515,15 @@ version = "0.3.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f2696e8a945f658fd14dc3b87242e6b80cd0f36ff04ea560fa39082368847946"
 
+[[package]]
+name = "des"
+version = "0.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ffdd80ce8ce993de27e9f063a444a4d53ce8e8db4c1f00cc03af5ad5a9867a1e"
+dependencies = [
+ "cipher",
+]
+
 [[package]]
 name = "digest"
 version = "0.10.7"
@@ -850,6 +877,15 @@ version = "0.4.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
 
+[[package]]
+name = "hmac"
+version = "0.12.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e"
+dependencies = [
+ "digest",
+]
+
 [[package]]
 name = "http"
 version = "0.2.9"
@@ -982,6 +1018,7 @@ version = "0.1.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a0c10553d664a4d0bcff9f4215d0aac67a639cc68ef660840afe309b807bc9f5"
 dependencies = [
+ "block-padding",
  "generic-array",
 ]
 
@@ -1312,6 +1349,23 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39"
 
+[[package]]
+name = "p12"
+version = "0.6.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d4873306de53fe82e7e484df31e1e947d61514b6ea2ed6cd7b45d63006fd9224"
+dependencies = [
+ "cbc",
+ "cipher",
+ "des",
+ "getrandom",
+ "hmac",
+ "lazy_static",
+ "rc2",
+ "sha1",
+ "yasna",
+]
+
 [[package]]
 name = "parking_lot"
 version = "0.12.1"
@@ -1540,13 +1594,17 @@ dependencies = [
  "lazy_static",
  "notify",
  "openssl",
+ "p12",
  "rand",
+ "rustls-native-certs",
+ "rustls-pemfile",
  "serde",
  "sha2",
  "snowstorm",
  "socket2 0.4.9",
  "tokio",
  "tokio-native-tls",
+ "tokio-rustls",
  "tokio-tungstenite",
  "tokio-util",
  "toml",
@@ -1556,6 +1614,15 @@ dependencies = [
  "vergen",
 ]
 
+[[package]]
+name = "rc2"
+version = "0.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "62c64daa8e9438b84aaae55010a93f396f8e60e3911590fcba770d04643fc1dd"
+dependencies = [
+ "cipher",
+]
+
 [[package]]
 name = "redox_syscall"
 version = "0.3.5"
@@ -1609,6 +1676,20 @@ version = "0.7.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "dbb5fb1acd8a1a18b3dd5be62d25485eb770e05afb408a9627d14d451bae12da"
 
+[[package]]
+name = "ring"
+version = "0.17.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "688c63d65483050968b2a8937f7995f443e27041a0f7700aa59b0822aedebb74"
+dependencies = [
+ "cc",
+ "getrandom",
+ "libc",
+ "spin",
+ "untrusted",
+ "windows-sys 0.48.0",
+]
+
 [[package]]
 name = "rustc-demangle"
 version = "0.1.23"
@@ -1637,6 +1718,60 @@ dependencies = [
  "windows-sys 0.48.0",
 ]
 
+[[package]]
+name = "rustls"
+version = "0.22.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e87c9956bd9807afa1f77e0f7594af32566e830e088a5576d27c5b6f30f49d41"
+dependencies = [
+ "log",
+ "ring",
+ "rustls-pki-types",
+ "rustls-webpki",
+ "subtle",
+ "zeroize",
+]
+
+[[package]]
+name = "rustls-native-certs"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8f1fb85efa936c42c6d5fc28d2629bb51e4b2f4b8a5211e297d599cc5a093792"
+dependencies = [
+ "openssl-probe",
+ "rustls-pemfile",
+ "rustls-pki-types",
+ "schannel",
+ "security-framework",
+]
+
+[[package]]
+name = "rustls-pemfile"
+version = "2.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "35e4980fa29e4c4b212ffb3db068a564cbf560e51d3944b7c88bd8bf5bec64f4"
+dependencies = [
+ "base64 0.21.4",
+ "rustls-pki-types",
+]
+
+[[package]]
+name = "rustls-pki-types"
+version = "1.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "048a63e5b3ac996d78d402940b5fa47973d2d080c6c6fffa1d0f19c4445310b7"
+
+[[package]]
+name = "rustls-webpki"
+version = "0.102.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "faaa0a62740bedb9b2ef5afa303da42764c012f743917351dc9a237ea1663610"
+dependencies = [
+ "ring",
+ "rustls-pki-types",
+ "untrusted",
+]
+
 [[package]]
 name = "rustversion"
 version = "1.0.14"
@@ -1838,6 +1973,12 @@ dependencies = [
  "windows-sys 0.48.0",
 ]
 
+[[package]]
+name = "spin"
+version = "0.9.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67"
+
 [[package]]
 name = "strsim"
 version = "0.10.0"
@@ -2030,6 +2171,17 @@ dependencies = [
  "tokio",
 ]
 
+[[package]]
+name = "tokio-rustls"
+version = "0.25.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "775e0c0f0adb3a2f22a00c4745d728b479985fc15ee7ca6a2608388c5569860f"
+dependencies = [
+ "rustls",
+ "rustls-pki-types",
+ "tokio",
+]
+
 [[package]]
 name = "tokio-stream"
 version = "0.1.14"
@@ -2261,6 +2413,12 @@ dependencies = [
  "subtle",
 ]
 
+[[package]]
+name = "untrusted"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1"
+
 [[package]]
 name = "url"
 version = "2.4.1"
@@ -2501,6 +2659,12 @@ version = "0.48.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538"
 
+[[package]]
+name = "yasna"
+version = "0.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e17bb3549cc1321ae1296b9cdc2698e2b6cb1992adfa19a8c72e5b7a738f44cd"
+
 [[package]]
 name = "zeroize"
 version = "1.6.0"

Cargo.toml

@@ -11,18 +11,48 @@ build = "build.rs"
 include = ["src/**/*", "LICENSE", "README.md", "build.rs"]
 
 [features]
-default = ["server", "client", "tls", "noise", "websocket", "hot-reload"]
+default = [
+    "server",
+    "client",
+    "native-tls",
+    "noise",
+    "websocket-native-tls",
+    "hot-reload",
+]
 
 # Run as a server
 server = []
 # Run as a client
 client = []
+
 # TLS support
-tls = ["tokio-native-tls"]
+native-tls = ["tokio-native-tls"]
+rustls = [
+    "tokio-rustls",
+    "rustls-pemfile",
+    "rustls-native-certs",
+    "p12",
+]
+
 # Noise support
 noise = ["snowstorm", "base64"]
+
 # Websocket support
-websocket = ["tokio-tungstenite", "tokio-util", "futures-core", "futures-sink", "tls"]
+websocket-native-tls = [
+    "tokio-tungstenite",
+    "tokio-util",
+    "futures-core",
+    "futures-sink",
+    "native-tls",
+]
+websocket-rustls = [
+    "tokio-tungstenite",
+    "tokio-util",
+    "futures-core",
+    "futures-sink",
+    "rustls",
+]
+
 # Configuration hot-reload support
 hot-reload = ["notify"]
 
@@ -67,27 +97,42 @@ hex = "0.4"
 rand = "0.8"
 backoff = { version = "0.4", features = ["tokio"] }
 tracing = "0.1"
-tracing-subscriber = { version="0.3", features=["env-filter"] }
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 socket2 = { version = "0.4", features = ["all"] }
 fdlimit = "0.2"
-tokio-native-tls = { version = "0.3", optional = true }
 async-trait = "0.1"
-snowstorm = { version = "0.4", optional = true, features = ["stream"], default-features = false }
+snowstorm = { version = "0.4", optional = true, features = [
+    "stream",
+], default-features = false }
 base64 = { version = "0.13", optional = true }
 notify = { version = "5.0.0-pre.13", optional = true }
-console-subscriber = { version = "0.1", optional = true, features = ["parking_lot"] }
+console-subscriber = { version = "0.1", optional = true, features = [
+    "parking_lot",
+] }
 atty = "0.2"
-async-http-proxy = { version = "1.2", features = ["runtime-tokio", "basic-auth"] }
+async-http-proxy = { version = "1.2", features = [
+    "runtime-tokio",
+    "basic-auth",
+] }
 async-socks5 = "0.5"
 url = { version = "2.2", features = ["serde"] }
-tokio-tungstenite = { version="0.20.1", optional = true}
-tokio-util = { version="0.7.9", optional = true, features = ["io"] }
-futures-core = { version="0.3.28", optional = true  }
-futures-sink = { version="0.3.28", optional = true }
+tokio-tungstenite = { version = "0.20.1", optional = true }
+tokio-util = { version = "0.7.9", optional = true, features = ["io"] }
+futures-core = { version = "0.3.28", optional = true }
+futures-sink = { version = "0.3.28", optional = true }
+tokio-native-tls = { version = "0.3", optional = true }
+tokio-rustls = { version = "0.25", optional = true }
+rustls-native-certs = { version = "0.7", optional = true }
+rustls-pemfile = { version = "2.0", optional = true }
+p12 = { version = "0.6.3", optional = true }
 
 [target.'cfg(target_env = "musl")'.dependencies]
 openssl = { version = "0.10", features = ["vendored"] }
 
 [build-dependencies]
-vergen = { version = "7.4.2", default-features = false, features = ["build", "git", "cargo"] }
+vergen = { version = "7.4.2", default-features = false, features = [
+    "build",
+    "git",
+    "cargo",
+] }
 anyhow = "1.0"

examples/tls/create_self_signed_cert.sh

@@ -56,7 +56,8 @@ openssl x509 -req \
     -sha256 -extfile cert.conf
 
 # create pkcs12
-openssl pkcs12 -export -out identity.pfx -inkey server.key -in server.crt -certfile rootCA.crt -passout pass:1234
+openssl pkcs12 -export -out identity.pfx -inkey server.key -in server.crt -certfile rootCA.crt \
+    -passout pass:1234 -keypbe PBE-SHA1-3DES -certpbe PBE-SHA1-3DES
 
 # clean up
-rm server.csr csr.conf cert.conf
+rm server.csr csr.conf cert.conf

examples/tls/rootCA.crt

@@ -1,20 +1,20 @@
 -----BEGIN CERTIFICATE-----
-MIIDTzCCAjegAwIBAgIUT2Hjb+eORMuX0zIwClSygNTJiSQwDQYJKoZIhvcNAQEL
+MIIDTzCCAjegAwIBAgIUHPYndZflmbDV/30C+BHQSiNvUTQwDQYJKoZIhvcNAQEL
 BQAwNzEQMA4GA1UEAwwHTXlPd25DQTELMAkGA1UEBhMCVVMxFjAUBgNVBAcMDVNh
-biBGcmFuc2lzY28wHhcNMjMwMzA3MTIzOTM5WhcNMjQwMjI2MTIzOTM5WjA3MRAw
+biBGcmFuc2lzY28wHhcNMjQwMjE1MDUwNDQ5WhcNMjUwMjA1MDUwNDQ5WjA3MRAw
 DgYDVQQDDAdNeU93bkNBMQswCQYDVQQGEwJVUzEWMBQGA1UEBwwNU2FuIEZyYW5z
-aXNjbzCCASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBAL4hFcu/+GeSQRR0
-XniadepJtCp3juIaHaYLMIsKg4fUSOiVlOCJU27wYa6xaYOcjSKpv7tmZ7YwFBwO
-dGdlcqAFD1nj+JCsHQAJKRIYWY6UklrQb0rd+67HXF03cN4sPGiAKXy52jaPYJIS
-oz5w8mfcz66b3q6fYmefyjwvqBl5nJApiWzBEtLPDKhmT6ST3VuQLdmYNEmL3lL9
-wVJu3R1L7gnzoUFdHyeOpAoALFAI8zfezI8IJsDLLdVfKZNZYm0PDB98ldlBQ2wf
-uXFTzuVHeifBFcUxhV5/U9c3Fp7UnuMD7/RAcABBE8aW6wFl246WjTk4v6r0QYgZ
-49BrnGMCAwEAAaNTMFEwHQYDVR0OBBYEFIwCXoKvHjF6mWhgNLwSEktXT9S/MB8G
-A1UdIwQYMBaAFIwCXoKvHjF6mWhgNLwSEktXT9S/MA8GA1UdEwEB/wQFMAMBAf8w
-DQYJKoZIhvcNAQELBQADggEBAIlSJqo9QJUZTE1SzafqihkSXBuLAKMNq+Box02o
-2tticlBV3BVpNZ4SbOs8oYN/Hmr2cDSmgbf4ZB1BqExarsrLnFuIrM4XWVzuFHSt
-oMSlE/OE6cO0wzqUlihmUfx2azuXKPLotAObD6fwNbUb03YxTpNrEqFxIjYn6g56
-Mp1Eo/Na2ptr41Nin2gHsynPOWdPhpBqBxnWMFz1pfZ7TB1h92DVqFN92fMzgvAT
-oJdTGl9hFTcS4XrYwOhhITNGn7oM9uTFpTd/IZbjAakcAnLcwRumthD32YJPpXqV
-JC2zJNBvEbQ4hdvZu3eNx5J8GU8wiMoJgYNy4zNMbM3qM+E=
+aXNjbzCCASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBAKIx0LdgvDrXGoGw
+XJ9s3Y+nr34NMPPLTbo/C2Yj1pD4mxZKK7d1VuwuBNM1h/WQLhA9+x4ZcKYZ1S1g
+3BRMuAdm/ZJyeeI1QDRqUlZD16ehPnY0Zy9sZX7oMKVS0m7l8zDv4nvDp9prC5yf
+8eoI7zoAWiMv/xPacYXFTAJbUb0VgovFyf3rzgIzs/NBF675FxrQtbhM2j4DdMkJ
+9UwRi+qmqtH/Z/Ddy4oMkPflEgKSgDEidmqa552CRExO3c+1ZbMEzq8iOUZ3Vb+g
+enfo0SwQUxQ9PEUOAd13siEXs51jZ7JqNmj1d/lEIbAuX8znWDqLYz9FUN4QNsim
+8Q/trBcCAwEAAaNTMFEwHQYDVR0OBBYEFP7eOqvUgs8/LOMonEZ6ubRaLkQMMB8G
+A1UdIwQYMBaAFP7eOqvUgs8/LOMonEZ6ubRaLkQMMA8GA1UdEwEB/wQFMAMBAf8w
+DQYJKoZIhvcNAQELBQADggEBABfLdbsbchr8Ep4mCv75ojWe11Mdd3Eg8EOePukC
+w918zqU6dZMmbnLtoXFk6QgFZnvD5MpmU4/d/BmvL9+CJJ9mJPwR2Vb/rIOPXV13
++kjHo/NwNbw5TdmPMbneyCjMdxRqmYKGoWYwbsI09YCK5Cb0J2fYmMrcACSVIUvz
+WC7CPPwTA3zvzf9xab+naoE1dbThRDGvVPXEFFOSMIXC0UzCvG0Lj3NTyXyu4XJ0
+TUcQUlnptLSejb+uh/5MSqwnEoc1dm2mW/oij1Gqg29+6WNw6wPv/cnC7VvlY4Eu
+CR9tvTjMNb7G6VRok9W0HJec6dNf3FJJ1pVzVL8bKI19G54=
 -----END CERTIFICATE-----

src/client.rs

@@ -8,8 +8,9 @@ use crate::protocol::{
 };
 use crate::transport::{AddrMaybeCached, SocketOpts, TcpTransport, Transport};
 use anyhow::{anyhow, bail, Context, Result};
+use backoff::backoff::Backoff;
+use backoff::future::retry_notify;
 use backoff::ExponentialBackoff;
-use backoff::{backoff::Backoff, future::retry_notify};
 use bytes::{Bytes, BytesMut};
 use std::collections::HashMap;
 use std::net::SocketAddr;
@@ -22,9 +23,9 @@ use tracing::{debug, error, info, instrument, trace, warn, Instrument, Span};
 
 #[cfg(feature = "noise")]
 use crate::transport::NoiseTransport;
-#[cfg(feature = "tls")]
+#[cfg(any(feature = "native-tls", feature = "rustls"))]
 use crate::transport::TlsTransport;
-#[cfg(feature = "websocket")]
+#[cfg(any(feature = "websocket-native-tls", feature = "websocket-rustls"))]
 use crate::transport::WebsocketTransport;
 
 use crate::constants::{run_control_chan_backoff, UDP_BUFFER_SIZE, UDP_SENDQ_SIZE, UDP_TIMEOUT};
@@ -47,13 +48,13 @@ pub async fn run_client(
             client.run(shutdown_rx, update_rx).await
         }
         TransportType::Tls => {
-            #[cfg(feature = "tls")]
+            #[cfg(any(feature = "native-tls", feature = "rustls"))]
             {
                 let mut client = Client::<TlsTransport>::from(config).await?;
                 client.run(shutdown_rx, update_rx).await
             }
-            #[cfg(not(feature = "tls"))]
-            crate::helper::feature_not_compile("tls")
+            #[cfg(not(any(feature = "native-tls", feature = "rustls")))]
+            crate::helper::feature_neither_compile("native-tls", "rustls")
         }
         TransportType::Noise => {
             #[cfg(feature = "noise")]
@@ -65,13 +66,13 @@ pub async fn run_client(
             crate::helper::feature_not_compile("noise")
         }
         TransportType::Websocket => {
-            #[cfg(feature = "websocket")]
+            #[cfg(any(feature = "websocket-native-tls", feature = "websocket-rustls"))]
             {
                 let mut client = Client::<WebsocketTransport>::from(config).await?;
                 client.run(shutdown_rx, update_rx).await
             }
-            #[cfg(not(feature = "websocket"))]
-            crate::helper::feature_not_compile("websocket")
+            #[cfg(not(any(feature = "websocket-native-tls", feature = "websocket-rustls")))]
+            crate::helper::feature_neither_compile("websocket-native-tls", "websocket-rustls")
         }
     }
 }

src/helper.rs

@@ -43,6 +43,14 @@ pub fn feature_not_compile(feature: &str) -> ! {
     )
 }
 
+#[allow(dead_code)]
+pub fn feature_neither_compile(feature1: &str, feature2: &str) -> ! {
+    panic!(
+        "Neither of the feature '{}' or '{}' is compiled in this binary. Please re-compile rathole",
+        feature1, feature2
+    )
+}
+
 pub async fn to_socket_addr<A: ToSocketAddrs>(addr: A) -> Result<SocketAddr> {
     lookup_host(addr)
         .await?

src/lib.rs

@@ -83,7 +83,7 @@ pub async fn run(args: Cli, shutdown_rx: broadcast::Receiver<bool>) -> Result<()
                 if let Some((i, _)) = last_instance {
                     info!("General configuration change detected. Restarting...");
                     shutdown_tx.send(true)?;
-                    i.await?;
+                    i.await??;
                 }
 
                 debug!("{:?}", config);
@@ -119,8 +119,8 @@ async fn run_instance(
     args: Cli,
     shutdown_rx: broadcast::Receiver<bool>,
     service_update: mpsc::Receiver<ConfigChange>,
-) {
-    let ret: Result<()> = match determine_run_mode(&config, &args) {
+) -> Result<()> {
+    match determine_run_mode(&config, &args) {
         RunMode::Undetermine => panic!("Cannot determine running as a server or a client"),
         RunMode::Client => {
             #[cfg(not(feature = "client"))]
@@ -134,8 +134,7 @@ async fn run_instance(
             #[cfg(feature = "server")]
             run_server(config, shutdown_rx, service_update).await
         }
-    };
-    ret.unwrap();
+    }
 }
 
 #[derive(PartialEq, Eq, Debug)]

src/server.rs

@@ -25,9 +25,9 @@ use tracing::{debug, error, info, info_span, instrument, warn, Instrument, Span}
 
 #[cfg(feature = "noise")]
 use crate::transport::NoiseTransport;
-#[cfg(feature = "tls")]
+#[cfg(any(feature = "native-tls", feature = "rustls"))]
 use crate::transport::TlsTransport;
-#[cfg(feature = "websocket")]
+#[cfg(any(feature = "websocket-native-tls", feature = "websocket-rustls"))]
 use crate::transport::WebsocketTransport;
 
 type ServiceDigest = protocol::Digest; // SHA256 of a service name
@@ -57,13 +57,13 @@ pub async fn run_server(
             server.run(shutdown_rx, update_rx).await?;
         }
         TransportType::Tls => {
-            #[cfg(feature = "tls")]
+            #[cfg(any(feature = "native-tls", feature = "rustls"))]
             {
                 let mut server = Server::<TlsTransport>::from(config).await?;
                 server.run(shutdown_rx, update_rx).await?;
             }
-            #[cfg(not(feature = "tls"))]
-            crate::helper::feature_not_compile("tls")
+            #[cfg(not(any(feature = "native-tls", feature = "rustls")))]
+            crate::helper::feature_neither_compile("native-tls", "rustls")
         }
         TransportType::Noise => {
             #[cfg(feature = "noise")]
@@ -75,13 +75,13 @@ pub async fn run_server(
             crate::helper::feature_not_compile("noise")
         }
         TransportType::Websocket => {
-            #[cfg(feature = "websocket")]
+            #[cfg(any(feature = "websocket-native-tls", feature = "websocket-rustls"))]
             {
                 let mut server = Server::<WebsocketTransport>::from(config).await?;
                 server.run(shutdown_rx, update_rx).await?;
             }
-            #[cfg(not(feature = "websocket"))]
-            crate::helper::feature_not_compile("websocket")
+            #[cfg(not(any(feature = "websocket-native-tls", feature = "websocket-rustls")))]
+            crate::helper::feature_neither_compile("websocket-native-tls", "websocket-rustls")
         }
     }
 

src/transport/mod.rs

@@ -69,19 +69,30 @@ pub trait Transport: Debug + Send + Sync {
 
 mod tcp;
 pub use tcp::TcpTransport;
-#[cfg(feature = "tls")]
-mod tls;
-#[cfg(feature = "tls")]
-pub use tls::TlsTransport;
+
+#[cfg(all(feature = "native-tls", feature = "rustls"))]
+compile_error!("Only one of `native-tls` and `rustls` can be enabled");
+
+#[cfg(feature = "native-tls")]
+mod native_tls;
+#[cfg(feature = "native-tls")]
+use native_tls as tls;
+#[cfg(feature = "rustls")]
+mod rustls;
+#[cfg(feature = "rustls")]
+use rustls as tls;
+
+#[cfg(any(feature = "native-tls", feature = "rustls"))]
+pub(crate) use tls::TlsTransport;
 
 #[cfg(feature = "noise")]
 mod noise;
 #[cfg(feature = "noise")]
 pub use noise::NoiseTransport;
 
-#[cfg(feature = "websocket")]
+#[cfg(any(feature = "websocket-native-tls", feature = "websocket-rustls"))]
 mod websocket;
-#[cfg(feature = "websocket")]
+#[cfg(any(feature = "websocket-native-tls", feature = "websocket-rustls"))]
 pub use websocket::WebsocketTransport;
 
 #[derive(Debug, Clone, Copy)]

src/transport/tls.rs → src/transport/native_tls.rs

@@ -1,14 +1,14 @@
-use std::net::SocketAddr;
-
-use super::{AddrMaybeCached, SocketOpts, TcpTransport, Transport};
 use crate::config::{TlsConfig, TransportConfig};
 use crate::helper::host_port_pair;
+use crate::transport::{AddrMaybeCached, SocketOpts, TcpTransport, Transport};
 use anyhow::{anyhow, Context, Result};
 use async_trait::async_trait;
 use std::fs;
+use std::net::SocketAddr;
 use tokio::net::{TcpListener, TcpStream, ToSocketAddrs};
 use tokio_native_tls::native_tls::{self, Certificate, Identity};
-use tokio_native_tls::{TlsAcceptor, TlsConnector, TlsStream};
+pub(crate) use tokio_native_tls::TlsStream;
+use tokio_native_tls::{TlsAcceptor, TlsConnector};
 
 #[derive(Debug)]
 pub struct TlsTransport {
@@ -109,3 +109,8 @@ impl Transport for TlsTransport {
             .await?)
     }
 }
+
+#[cfg(feature = "websocket-native-tls")]
+pub(crate) fn get_tcpstream(s: &TlsStream<TcpStream>) -> &TcpStream {
+    s.get_ref().get_ref().get_ref()
+}

src/transport/rustls.rs

@@ -0,0 +1,156 @@
+use crate::config::{TlsConfig, TransportConfig};
+use crate::helper::host_port_pair;
+use crate::transport::{AddrMaybeCached, SocketOpts, TcpTransport, Transport};
+use std::fmt::Debug;
+use std::fs;
+use std::net::SocketAddr;
+use std::sync::Arc;
+use tokio::net::{TcpListener, TcpStream, ToSocketAddrs};
+use tokio_rustls::rustls::pki_types::{CertificateDer, PrivatePkcs8KeyDer, ServerName};
+
+use anyhow::{anyhow, Context, Result};
+use async_trait::async_trait;
+use p12::PFX;
+use tokio_rustls::rustls::{ClientConfig, RootCertStore, ServerConfig};
+pub(crate) use tokio_rustls::TlsStream;
+use tokio_rustls::{TlsAcceptor, TlsConnector};
+
+pub struct TlsTransport {
+    tcp: TcpTransport,
+    config: TlsConfig,
+    connector: Option<TlsConnector>,
+    tls_acceptor: Option<TlsAcceptor>,
+}
+
+// workaround for TlsConnector and TlsAcceptor not implementing Debug
+impl Debug for TlsTransport {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("TlsTransport")
+            .field("tcp", &self.tcp)
+            .field("config", &self.config)
+            .finish()
+    }
+}
+
+fn load_server_config(config: &TlsConfig) -> Result<Option<ServerConfig>> {
+    if let Some(pkcs12_path) = config.pkcs12.as_ref() {
+        let buf = fs::read(pkcs12_path)?;
+        let pfx = PFX::parse(buf.as_slice())?;
+        let pass = config.pkcs12_password.as_ref().unwrap();
+
+        let certs = pfx.cert_bags(pass)?;
+        let keys = pfx.key_bags(pass)?;
+
+        let chain: Vec<CertificateDer> = certs.into_iter().map(CertificateDer::from).collect();
+        let key = PrivatePkcs8KeyDer::from(keys.into_iter().next().unwrap());
+
+        Ok(Some(
+            ServerConfig::builder()
+                .with_no_client_auth()
+                .with_single_cert(chain, key.into())?,
+        ))
+    } else {
+        Ok(None)
+    }
+}
+
+fn load_client_config(config: &TlsConfig) -> Result<Option<ClientConfig>> {
+    let cert = if let Some(path) = config.trusted_root.as_ref() {
+        rustls_pemfile::certs(&mut std::io::BufReader::new(fs::File::open(path).unwrap()))
+            .map(|cert| cert.unwrap())
+            .next()
+            .with_context(|| "Failed to read certificate")?
+    } else {
+        // read from native
+        match rustls_native_certs::load_native_certs() {
+            Ok(certs) => certs.into_iter().next().unwrap(),
+            Err(e) => {
+                eprintln!("Failed to load native certs: {}", e);
+                return Ok(None);
+            }
+        }
+    };
+
+    let mut root_certs = RootCertStore::empty();
+    root_certs.add(cert).unwrap();
+
+    Ok(Some(
+        ClientConfig::builder()
+            .with_root_certificates(root_certs)
+            .with_no_client_auth(),
+    ))
+}
+
+#[async_trait]
+impl Transport for TlsTransport {
+    type Acceptor = TcpListener;
+    type RawStream = TcpStream;
+    type Stream = TlsStream<TcpStream>;
+
+    fn new(config: &TransportConfig) -> Result<Self> {
+        let tcp = TcpTransport::new(config)?;
+        let config = config
+            .tls
+            .as_ref()
+            .ok_or_else(|| anyhow!("Missing tls config"))?;
+
+        let connector = load_client_config(config)
+            .unwrap()
+            .map(|c| Arc::new(c).into());
+        let tls_acceptor = load_server_config(config)
+            .unwrap()
+            .map(|c| Arc::new(c).into());
+
+        Ok(TlsTransport {
+            tcp,
+            config: config.clone(),
+            connector,
+            tls_acceptor,
+        })
+    }
+
+    fn hint(conn: &Self::Stream, opt: SocketOpts) {
+        opt.apply(conn.get_ref().0);
+    }
+
+    async fn bind<A: ToSocketAddrs + Send + Sync>(&self, addr: A) -> Result<Self::Acceptor> {
+        let l = TcpListener::bind(addr)
+            .await
+            .with_context(|| "Failed to create tcp listener")?;
+        Ok(l)
+    }
+
+    async fn accept(&self, a: &Self::Acceptor) -> Result<(Self::RawStream, SocketAddr)> {
+        self.tcp
+            .accept(a)
+            .await
+            .with_context(|| "Failed to accept TCP connection")
+    }
+
+    async fn handshake(&self, conn: Self::RawStream) -> Result<Self::Stream> {
+        let conn = self.tls_acceptor.as_ref().unwrap().accept(conn).await?;
+        Ok(tokio_rustls::TlsStream::Server(conn))
+    }
+
+    async fn connect(&self, addr: &AddrMaybeCached) -> Result<Self::Stream> {
+        let conn = self.tcp.connect(addr).await?;
+
+        let connector = self.connector.as_ref().unwrap();
+
+        let host_name = self
+            .config
+            .hostname
+            .as_deref()
+            .unwrap_or(host_port_pair(&addr.addr)?.0);
+
+        Ok(tokio_rustls::TlsStream::Client(
+            connector
+                .connect(ServerName::try_from(host_name)?.to_owned(), conn)
+                .await?,
+        ))
+    }
+}
+
+pub(crate) fn get_tcpstream(s: &TlsStream<TcpStream>) -> &TcpStream {
+    &s.get_ref().0
+}

src/transport/websocket.rs

@@ -13,10 +13,14 @@ use futures_core::stream::Stream;
 use futures_sink::Sink;
 use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf};
 use tokio::net::{TcpListener, TcpStream, ToSocketAddrs};
-use tokio_native_tls::TlsStream;
-use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
-use tokio_tungstenite::{accept_async_with_config, client_async_with_config};
-use tokio_tungstenite::{tungstenite::protocol::Message, WebSocketStream};
+
+#[cfg(any(feature = "native-tls", feature = "rustls"))]
+use super::tls::get_tcpstream;
+#[cfg(any(feature = "native-tls", feature = "rustls"))]
+use super::tls::TlsStream;
+
+use tokio_tungstenite::tungstenite::protocol::{Message, WebSocketConfig};
+use tokio_tungstenite::{accept_async_with_config, client_async_with_config, WebSocketStream};
 use tokio_util::io::StreamReader;
 use url::Url;
 
@@ -30,7 +34,7 @@ impl TransportStream {
     fn get_tcpstream(&self) -> &TcpStream {
         match self {
             TransportStream::Insecure(s) => s,
-            TransportStream::Secure(s) => s.get_ref().get_ref().get_ref(),
+            TransportStream::Secure(s) => get_tcpstream(s),
         }
     }
 }

tests/integration_test.rs

@@ -1,4 +1,4 @@
-use anyhow::Result;
+use anyhow::{Ok, Result};
 use common::{run_rathole_client, PING, PONG};
 use rand::Rng;
 use std::time::Duration;
@@ -57,17 +57,17 @@ async fn tcp() -> Result<()> {
     test("tests/for_tcp/tcp_transport.toml", Type::Tcp).await?;
     // FIXME: Self-signed certificate on Mac requires mannual interference. Disable CI for now
     #[cfg(not(target_os = "macos"))]
-    #[cfg(feature="tls")]
+    #[cfg(any(feature = "native-tls", feature = "rustls"))]
     test("tests/for_tcp/tls_transport.toml", Type::Tcp).await?;
 
-    #[cfg(feature="noise")]
+    #[cfg(feature = "noise")]
     test("tests/for_tcp/noise_transport.toml", Type::Tcp).await?;
 
-    #[cfg(feature="websocket")]
+    #[cfg(any(feature = "websocket-native-tls", feature = "websocket-rustls"))]
     test("tests/for_tcp/websocket_transport.toml", Type::Tcp).await?;
 
     #[cfg(not(target_os = "macos"))]
-    #[cfg(feature="websocket")]
+    #[cfg(any(feature = "websocket-native-tls", feature = "websocket-rustls"))]
     test("tests/for_tcp/websocket_tls_transport.toml", Type::Tcp).await?;
 
     Ok(())
@@ -94,17 +94,17 @@ async fn udp() -> Result<()> {
     test("tests/for_udp/tcp_transport.toml", Type::Udp).await?;
     // See above
     #[cfg(not(target_os = "macos"))]
-    #[cfg(feature="tls")]
+    #[cfg(any(feature = "native-tls", feature = "rustls"))]
     test("tests/for_udp/tls_transport.toml", Type::Udp).await?;
 
-    #[cfg(feature="noise")]
+    #[cfg(feature = "noise")]
     test("tests/for_udp/noise_transport.toml", Type::Udp).await?;
 
-    #[cfg(feature="websocket")]
+    #[cfg(any(feature = "websocket-native-tls", feature = "websocket-rustls"))]
     test("tests/for_udp/websocket_transport.toml", Type::Udp).await?;
 
     #[cfg(not(target_os = "macos"))]
-    #[cfg(feature="websocket")]
+    #[cfg(any(feature = "websocket-native-tls", feature = "websocket-rustls"))]
     test("tests/for_udp/websocket_tls_transport.toml", Type::Udp).await?;
 
     Ok(())
@@ -112,6 +112,11 @@ async fn udp() -> Result<()> {
 
 #[instrument]
 async fn test(config_path: &'static str, t: Type) -> Result<()> {
+    if cfg!(not(all(feature = "client", feature = "server"))) {
+        // Skip the test if the client or the server is not enabled
+        return Ok(());
+    }
+
     let (client_shutdown_tx, client_shutdown_rx) = broadcast::channel(1);
     let (server_shutdown_tx, server_shutdown_rx) = broadcast::channel(1);