gateway: gRPC proxy support

Author: kvinwang

gateway/src/main_service.rs

@@ -60,6 +60,7 @@ pub struct ProxyInner {
     notify_state_updated: Notify,
     auth_client: AuthClient,
     pub(crate) acceptor: RwLock<TlsAcceptor>,
+    pub(crate) h2_acceptor: RwLock<TlsAcceptor>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -141,15 +142,19 @@ impl ProxyInner {
             }
             false => None,
         };
-        let acceptor =
-            RwLock::new(create_acceptor(&config.proxy).context("Failed to create acceptor")?);
+        let acceptor = RwLock::new(
+            create_acceptor(&config.proxy, false).context("Failed to create acceptor")?,
+        );
+        let h2_acceptor =
+            RwLock::new(create_acceptor(&config.proxy, true).context("Failed to create acceptor")?);
         Ok(Self {
             config,
             state,
             notify_state_updated: Notify::new(),
             my_app_id,
             auth_client,
             acceptor,
+            h2_acceptor,
             certbot,
         })
     }

gateway/src/proxy.rs

@@ -65,6 +65,7 @@ struct DstInfo {
     app_id: String,
     port: u16,
     is_tls: bool,
+    is_h2: bool,
 }
 
 fn parse_destination(sni: &str, dotted_base_domain: &str) -> Result<DstInfo> {
@@ -83,22 +84,28 @@ fn parse_destination(sni: &str, dotted_base_domain: &str) -> Result<DstInfo> {
     let last_part = parts.next();
     let is_tls;
     let port;
+    let is_h2;
     match last_part {
         None => {
             is_tls = false;
+            is_h2 = false;
             port = None;
         }
         Some(last_part) => {
-            let port_str = match last_part.strip_suffix('s') {
-                None => {
-                    is_tls = false;
-                    last_part
-                }
-                Some(last_part) => {
-                    is_tls = true;
-                    last_part
-                }
+            let (port_str, has_g) = match last_part.strip_suffix('g') {
+                Some(without_g) => (without_g, true),
+                None => (last_part, false),
             };
+
+            let (port_str, has_s) = match port_str.strip_suffix('s') {
+                Some(without_s) => (without_s, true),
+                None => (port_str, false),
+            };
+            if has_g && has_s {
+                bail!("invalid sni format: `gs` is not allowed");
+            }
+            is_h2 = has_g;
+            is_tls = has_s;
             port = if port_str.is_empty() {
                 None
             } else {
@@ -114,6 +121,7 @@ fn parse_destination(sni: &str, dotted_base_domain: &str) -> Result<DstInfo> {
         app_id,
         port,
         is_tls,
+        is_h2,
     })
 }
 
@@ -138,7 +146,9 @@ async fn handle_connection(
         if dst.is_tls {
             tls_passthough::proxy_to_app(state, inbound, buffer, &dst.app_id, dst.port).await
         } else {
-            state.proxy(inbound, buffer, &dst.app_id, dst.port).await
+            state
+                .proxy(inbound, buffer, &dst.app_id, dst.port, dst.is_h2)
+                .await
         }
     } else {
         tls_passthough::proxy_with_sni(state, inbound, buffer, &sni).await

gateway/src/proxy/tls_terminate.rs

@@ -93,7 +93,7 @@ where
     }
 }
 
-pub(crate) fn create_acceptor(config: &ProxyConfig) -> Result<TlsAcceptor> {
+pub(crate) fn create_acceptor(config: &ProxyConfig, h2: bool) -> Result<TlsAcceptor> {
     let cert_pem = fs::read(&config.cert_chain).context("failed to read certificate")?;
     let key_pem = fs::read(&config.cert_key).context("failed to read private key")?;
     let certs = CertificateDer::pem_slice_iter(cert_pem.as_slice())
@@ -114,12 +114,16 @@ pub(crate) fn create_acceptor(config: &ProxyConfig) -> Result<TlsAcceptor> {
             TlsVersion::Tls13 => &TLS13,
         })
         .collect::<Vec<_>>();
-    let config = rustls::ServerConfig::builder_with_provider(Arc::new(provider))
+    let mut config = rustls::ServerConfig::builder_with_provider(Arc::new(provider))
         .with_protocol_versions(&supported_versions)
         .context("Failed to build TLS config")?
         .with_no_client_auth()
         .with_single_cert(certs, key)?;
 
+    if h2 {
+        config.alpn_protocols = vec![b"h2".to_vec()];
+    }
+
     let acceptor = TlsAcceptor::from(Arc::new(config));
 
     Ok(acceptor)
@@ -145,11 +149,16 @@ impl Proxy {
     /// Reload the TLS acceptor with fresh certificates
     pub fn reload_certificates(&self) -> Result<()> {
         info!("Reloading TLS certificates");
-        let new_acceptor = create_acceptor(&self.config.proxy)?;
-
         // Replace the acceptor with the new one
         if let Ok(mut acceptor) = self.acceptor.write() {
-            *acceptor = new_acceptor;
+            *acceptor = create_acceptor(&self.config.proxy, false)?;
+            info!("TLS certificates successfully reloaded");
+        } else {
+            bail!("Failed to acquire write lock for TLS acceptor");
+        }
+
+        if let Ok(mut acceptor) = self.h2_acceptor.write() {
+            *acceptor = create_acceptor(&self.config.proxy, true)?;
             info!("TLS certificates successfully reloaded");
         } else {
             bail!("Failed to acquire write lock for TLS acceptor");
@@ -163,11 +172,12 @@ impl Proxy {
         inbound: TcpStream,
         buffer: Vec<u8>,
         port: u16,
+        h2: bool,
     ) -> Result<()> {
         if port != 80 {
             bail!("Only port 80 is supported for this node");
         }
-        let stream = self.tls_accept(inbound, buffer).await?;
+        let stream = self.tls_accept(inbound, buffer, h2).await?;
         let io = TokioIo::new(stream);
 
         let service = service_fn(|req: Request<Incoming>| async move {
@@ -218,11 +228,12 @@ impl Proxy {
         inbound: TcpStream,
         buffer: Vec<u8>,
         port: u16,
+        h2: bool,
     ) -> Result<()> {
         if port != 80 {
             bail!("Only port 80 is supported for health checks");
         }
-        let stream = self.tls_accept(inbound, buffer).await?;
+        let stream = self.tls_accept(inbound, buffer, h2).await?;
 
         // Wrap the TLS stream with TokioIo to make it compatible with hyper 1.x
         let io = TokioIo::new(stream);
@@ -253,17 +264,24 @@ impl Proxy {
         &self,
         inbound: TcpStream,
         buffer: Vec<u8>,
+        h2: bool,
     ) -> Result<TlsStream<MergedStream>> {
         let stream = MergedStream {
             buffer,
             buffer_cursor: 0,
             inbound,
         };
-        let acceptor = self
-            .acceptor
-            .read()
-            .expect("Failed to acquire read lock for TLS acceptor")
-            .clone();
+        let acceptor = if h2 {
+            self.h2_acceptor
+                .read()
+                .expect("Failed to acquire read lock for TLS acceptor")
+                .clone()
+        } else {
+            self.acceptor
+                .read()
+                .expect("Failed to acquire read lock for TLS acceptor")
+                .clone()
+        };
         let tls_stream = timeout(
             self.config.proxy.timeouts.handshake,
             acceptor.accept(stream),
@@ -280,19 +298,20 @@ impl Proxy {
         buffer: Vec<u8>,
         app_id: &str,
         port: u16,
+        h2: bool,
     ) -> Result<()> {
         if app_id == "health" {
-            return self.handle_health_check(inbound, buffer, port).await;
+            return self.handle_health_check(inbound, buffer, port, h2).await;
         }
         if app_id == "gateway" {
-            return self.handle_this_node(inbound, buffer, port).await;
+            return self.handle_this_node(inbound, buffer, port, h2).await;
         }
         let addresses = self
             .lock()
             .select_top_n_hosts(app_id)
             .with_context(|| format!("app {app_id} not found"))?;
         debug!("selected top n hosts: {addresses:?}");
-        let tls_stream = self.tls_accept(inbound, buffer).await?;
+        let tls_stream = self.tls_accept(inbound, buffer, h2).await?;
         let (outbound, _counter) = timeout(
             self.config.proxy.timeouts.connect,
             connect_multiple_hosts(addresses, port),