diff --git a/Cargo.lock b/Cargo.lock index bd5994a8..b19500e2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -95,6 +95,23 @@ version = "1.0.70" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7de8ce5e0f9f8d88245311066a578d72b7af3e7088f32783804676302df237e4" +[[package]] +name = "arc-swap" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bddcadddf5e9015d310179a59bb28c4d4b9920ad0f11e8e14dbadf654890c9a6" + +[[package]] +name = "async-rustls" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93b21a03b7c21702a0110f9f8d228763a533570deb376119042dabf33c37a01a" +dependencies = [ + "futures-io", + "rustls", + "webpki", +] + [[package]] name = "async-stream" version = "0.3.5" @@ -145,6 +162,18 @@ dependencies = [ "num-traits", ] +[[package]] +name = "awaitdrop" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "771051cdc7eec2dc1b23fbf870bb7fbb89136fe374227c875e377f1eed99a429" +dependencies = [ + "futures", + "generational-arena", + "parking_lot", + "slotmap", +] + [[package]] name = "axum" version = "0.6.15" @@ -758,7 +787,7 @@ dependencies = [ "futures-sink", "nanorand", "pin-project", - "spin", + "spin 0.9.8", ] [[package]] @@ -890,6 +919,15 @@ dependencies = [ "slab", ] +[[package]] +name = "generational-arena" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877e94aff08e743b651baaea359664321055749b398adff8740a7399af7796e7" +dependencies = [ + "cfg-if", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -987,6 +1025,17 @@ dependencies = [ "digest", ] +[[package]] +name = "hostname" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c731c3e10504cc8ed35cfe2f1db4c9274c3d35fa486e3b31df46f068ef3e867" +dependencies = [ + "libc", + "match_cfg", + "winapi", +] + [[package]] name = "http" version = "0.2.9" @@ -1284,6 +1333,12 @@ version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "58093314a45e00c77d5c508f76e77c3396afbbc0d01506e7fae47b018bac2b1d" +[[package]] +name = "match_cfg" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffbee8634e0d45d258acb448e7eaab3fce7a0a467395d4d9f228e3c1f01fb2e4" + [[package]] name = "matchers" version = "0.1.0" @@ -1442,6 +1497,25 @@ version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e5ce46fe64a9d73be07dcbe690a38ce1b293be448fd8ce1e6c1b8062c9f72c6a" +[[package]] +name = "muxado" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e92b89ac3127251efde6f5a9586e5aae99468d06fcf9f133b377f58d5ed66446" +dependencies = [ + "async-trait", + "awaitdrop", + "bitflags", + "bytes", + "futures", + "pin-project", + "rand", + "thiserror", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "nanorand" version = "0.7.0" @@ -1469,6 +1543,37 @@ dependencies = [ "tempfile", ] +[[package]] +name = "ngrok" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98ce3514eec7338e2d4663e3efb4429e08d8f3662996be4b9585350e7d8ad728" +dependencies = [ + "arc-swap", + "async-rustls", + "async-trait", + "awaitdrop", + "axum", + "base64 0.13.1", + "bytes", + "futures", + "hostname", + "hyper", + "muxado", + "once_cell", + "parking_lot", + "regex", + "rustls-pemfile", + "serde", + "serde_json", + "thiserror", + "tokio", + "tokio-retry", + "tokio-util", + "tracing", + "windows-sys 0.45.0", +] + [[package]] name = "nix" version = "0.26.2" @@ -2140,6 +2245,21 @@ dependencies = [ "winreg", ] +[[package]] +name = "ring" +version = "0.16.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3053cf52e236a3ed746dfc745aa9cacf1b791d846bdaf412f60a8d7d6e17c8fc" +dependencies = [ + "cc", + "libc", + "once_cell", + "spin 0.5.2", + "untrusted", + "web-sys", + "winapi", +] + [[package]] name = "rust-embed" version = "6.6.1" @@ -2198,6 +2318,27 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "rustls" +version = "0.20.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fff78fc74d175294f4e83b28343315ffcfb114b156f0185e9741cb5570f50e2f" +dependencies = [ + "log", + "ring", + "sct", + "webpki", +] + +[[package]] +name = "rustls-pemfile" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d194b56d58803a43635bdc398cd17e383d6f71f9182b9a192c127ca42494a59b" +dependencies = [ + "base64 0.21.0", +] + [[package]] name = "rustversion" version = "1.0.12" @@ -2234,6 +2375,16 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" +[[package]] +name = "sct" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d53dcdb7c9f8158937a7981b48accfd39a43af418591a5d008c7b22b5e1b7ca4" +dependencies = [ + "ring", + "untrusted", +] + [[package]] name = "security-framework" version = "2.8.2" @@ -2400,6 +2551,15 @@ dependencies = [ "autocfg", ] +[[package]] +name = "slotmap" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1e08e261d0e8f5c43123b7adf3e4ca1690d655377ac93a03b2c9d3e98de1342" +dependencies = [ + "version_check", +] + [[package]] name = "smallvec" version = "1.10.0" @@ -2416,6 +2576,12 @@ dependencies = [ "winapi", ] +[[package]] +name = "spin" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e63cff320ae2c57904679ba7cb63280a3dc4613885beafb148ee7bf9aa9042d" + [[package]] name = "spin" version = "0.9.8" @@ -2619,6 +2785,7 @@ dependencies = [ "futures", "metrics", "metrics-exporter-prometheus", + "ngrok", "nohash-hasher", "opentelemetry", "opentelemetry-otlp", @@ -2797,6 +2964,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-retry" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f57eb36ecbe0fc510036adff84824dd3c24bb781e21bfa67b69d556aa85214f" +dependencies = [ + "pin-project", + "rand", + "tokio", +] + [[package]] name = "tokio-stream" version = "0.1.12" @@ -2816,6 +2994,7 @@ checksum = "5427d89453009325de0d8f342c9490009f76e999cb7672d77e46267448f7e6b2" dependencies = [ "bytes", "futures-core", + "futures-io", "futures-sink", "pin-project-lite", "tokio", @@ -3105,6 +3284,12 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" +[[package]] +name = "untrusted" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a156c684c91ea7d62626509bce3cb4e1d9ed5c4d978f7b4352658f96a4c26b4a" + [[package]] name = "url" version = "2.3.1" @@ -3301,6 +3486,16 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "webpki" +version = "0.22.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f095d78192e208183081cc07bc5515ef55216397af48b873e5edcd72637fa1bd" +dependencies = [ + "ring", + "untrusted", +] + [[package]] name = "which" version = "4.4.0" diff --git a/router/Cargo.toml b/router/Cargo.toml index 6503e1bd..722aadf7 100644 --- a/router/Cargo.toml +++ b/router/Cargo.toml @@ -40,6 +40,7 @@ tracing-opentelemetry = "0.18.0" tracing-subscriber = { version = "0.3.16", features = ["json", "env-filter"] } utoipa = { version = "3.0.1", features = ["axum_extras"] } utoipa-swagger-ui = { version = "3.0.2", features = ["axum"] } +ngrok = { version = "0.12.3", features = ["axum"] } [build-dependencies] vergen = { version = "8.0.0", features = ["build", "git", "gitcl"] } diff --git a/router/src/server.rs b/router/src/server.rs index 10c0ba3c..ab069698 100644 --- a/router/src/server.rs +++ b/router/src/server.rs @@ -1,5 +1,5 @@ -use crate::health::Health; /// HTTP Server logic +use crate::health::Health; use crate::infer::{InferError, InferResponse, InferStreamResponse}; use crate::validation::ValidationError; use crate::{ @@ -17,7 +17,10 @@ use axum_tracing_opentelemetry::opentelemetry_tracing_layer; use futures::stream::StreamExt; use futures::Stream; use metrics_exporter_prometheus::{Matcher, PrometheusBuilder, PrometheusHandle}; +use ngrok::config::TunnelBuilder; +use ngrok::tunnel::UrlTunnel; use std::convert::Infallible; +use std::env; use std::net::SocketAddr; use std::sync::atomic::AtomicBool; use std::sync::Arc; @@ -683,13 +686,45 @@ pub async fn run( .layer(opentelemetry_tracing_layer()) .layer(cors_layer); - // Run server - axum::Server::bind(&addr) - .serve(app.into_make_service()) - // Wait until all requests are finished to shut down - .with_graceful_shutdown(shutdown_signal()) - .await - .unwrap(); + if let Ok(ngrok_token) = env::var("NGROK_AUTHTOKEN") { + let _ = addr; + + let mut tunnel = ngrok::Session::builder() + .authtoken(ngrok_token) + .connect() + .await + .unwrap() + .http_endpoint(); + + if let Ok(domain) = env::var("NGROK_DOMAIN") { + tunnel = tunnel.domain(domain); + } + + if let (Ok(username), Ok(password)) = + (env::var("NGROK_USERNAME"), env::var("NGROK_PASSWORD")) + { + tunnel = tunnel.basic_auth(username, password); + } + + let listener = tunnel.listen().await.unwrap(); + + // Run server + tracing::info!("Ingress URL: {:?}", listener.url()); + axum::Server::builder(listener) + .serve(app.into_make_service()) + //Wait until all requests are finished to shut down + .with_graceful_shutdown(shutdown_signal()) + .await + .unwrap(); + } else { + // Run server + axum::Server::bind(&addr) + .serve(app.into_make_service()) + // Wait until all requests are finished to shut down + .with_graceful_shutdown(shutdown_signal()) + .await + .unwrap(); + } } /// Shutdown signal handler