From e850941f40fb150b7077a954886eb0e7bb9a1830 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 29 Jun 2020 22:54:05 -0400 Subject: [PATCH] Working proxy design --- Cargo.lock | 65 ++++++++++++++++++++++++++++ Cargo.toml | 1 + proxy_configs.toml | 6 +-- src/main.rs | 105 ++++++++++++++++++++++++++++++++------------- 4 files changed, 144 insertions(+), 33 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 82452be..5e1f55a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -175,6 +175,17 @@ dependencies = [ "num_cpus", ] +[[package]] +name = "getrandom" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7abc8dd8451921606d809ba32e95b6111925cd2906060d2dcc29c070220503eb" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + [[package]] name = "h2" version = "0.1.26" @@ -445,6 +456,12 @@ dependencies = [ "winapi 0.3.9", ] +[[package]] +name = "ppv-lite86" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "237a5ed80e274dbc66f86bd59c1e25edc039660be53194b5fe0a482e0f2612ea" + [[package]] name = "proc-macro2" version = "1.0.18" @@ -469,6 +486,47 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rand" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a6b1679d49b24bbfe0c803429aa1874472f50d9b363131f0e89fc356b544d03" +dependencies = [ + "getrandom", + "libc", + "rand_chacha", + "rand_core", + "rand_hc", +] + +[[package]] +name = "rand_chacha" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4c8ed856279c9737206bf725bf36935d8666ead7aa69b52be55af369d193402" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90bde5296fc891b0cef12a6d03ddccc162ce7b2aff54160af9338f8d40df6d19" +dependencies = [ + "getrandom", +] + +[[package]] +name = "rand_hc" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca3129af7b92a17112d59ad498c6f81eaf463253766b90396d39ea7a39d6613c" +dependencies = [ + "rand_core", +] + [[package]] name = "redox_syscall" version = "0.1.56" @@ -777,6 +835,7 @@ dependencies = [ "hyper", "hyper-reverse-proxy", "log", + "rand", "serde", "serde_derive", "serde_json", @@ -815,6 +874,12 @@ dependencies = [ "try-lock", ] +[[package]] +name = "wasi" +version = "0.9.0+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cccddf32554fecc6acb585f82a32a72e28b48f8c4c1883ddfeeeaa96f7d8e519" + [[package]] name = "winapi" version = "0.2.8" diff --git a/Cargo.toml b/Cargo.toml index 3dad6d2..07267c4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,6 +19,7 @@ futures = "0.1" hyper = "0.12" hyper-reverse-proxy = "0.4" log = "0.4.8" +rand = "0.7.3" serde = "1.0.104" serde_derive = "1.0.104" serde_json = "1.0.55" diff --git a/proxy_configs.toml b/proxy_configs.toml index c813adc..414d2e0 100644 --- a/proxy_configs.toml +++ b/proxy_configs.toml @@ -1,9 +1,9 @@ [[backends]] -voice = "foo" +voice = "trump" ip = "127.0.0.1" port = "12345" [[backends]] -voice = "bar" -ip = "1.1.1.1" +voice = "glow-ljs-single" +ip = "localhost" port = "12345" diff --git a/src/main.rs b/src/main.rs index 2cc38d9..cf3970d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,35 +2,33 @@ #[macro_use] extern crate log; #[macro_use] extern crate serde_derive; -use anyhow::{Result as AnyhowResult, Error}; +use anyhow::Result as AnyhowResult; use futures::future::{self, Future}; use hyper::Method; use hyper::rt::Stream; use hyper::server::conn::AddrStream; use hyper::service::{service_fn, make_service_fn}; use hyper::{Body, Request, Response, Server}; +use std::collections::HashMap; use std::env; use std::fmt::Debug; use std::fmt::Display; -use std::str::FromStr; use std::net::SocketAddr; -use std::collections::HashMap; -use hyper::header::HeaderValue; -use hyper::http::header::HeaderName; +use std::str::FromStr; +use std::sync::Arc; +use rand::seq::IteratorRandom; type BoxFut = Box, Error=hyper::Error> + Send>; +const ENV_DEFAULT_ROUTE : &'static str = "DEFAULT_ROUTE"; const ENV_PROXY_CONFIG_FILE : &'static str = "PROXY_CONFIG_FILE"; -const ENV_ROUTE_ONE : &'static str = "ROUTE_ONE"; -const ENV_ROUTE_TWO : &'static str = "ROUTE_TWO"; const ENV_RUST_LOG : &'static str = "RUST_LOG"; const ENV_SERVICE_PORT : &'static str = "SERVICE_PORT"; +const DEFAULT_DEFAULT_ROUTE : &'static str = "http://127.0.0.1:12345"; const DEFAULT_PROXY_CONFIG_FILE : &'static str = "proxy_configs.toml"; -const DEFAULT_RUST_LOG: &'static str = "debug"; +const DEFAULT_RUST_LOG: &'static str = "tokio_reactor=warn,hyper=info,debug"; const DEFAULT_SERVICE_PORT: u16 = 5555; -const ROUTE_ONE_DEFAULT : &'static str = "http://127.0.0.1:12345"; -const ROUTE_TWO_DEFAULT : &'static str = "http://127.0.0.1:3000"; #[derive(Deserialize, Debug, Clone)] struct ProxyConfigs { @@ -91,7 +89,7 @@ fn get_env_num(env_name: &str, default: T) -> AnyhowResult } } -fn debug_request(req: Request) -> BoxFut { +fn _debug_request(req: Request) -> BoxFut { let body_str = format!("{:?}", req); let response = Response::new(Body::from(body_str)); Box::new(future::ok(response)) @@ -102,8 +100,26 @@ struct RequestDetails { pub speaker: String, } -fn speak_proxy(req: Request, remote_addr: SocketAddr, endpoint: &'static str) -> BoxFut { - let mut headers = req.headers().clone(); +struct Router { + pub routes: HashMap, +} + +impl Router { + fn get_random_host(&self) -> Option { + let mut rng = rand::thread_rng(); + self.routes.values() + .choose(&mut rng) + .map(|url| url.to_string()) + } + + fn get_speaker_host(&self, voice: &str) -> Option { + self.routes.get(voice) + .map(|url| url.to_string()) + } +} + +fn speak_proxy(req: Request, remote_addr: SocketAddr, router: Arc, endpoint: &'static str) -> BoxFut { + let headers = req.headers().clone(); Box::new(req.into_body().concat2().map(move |b| { // Builds a BoxedFut to return let request_bytes = b.as_ref(); @@ -116,6 +132,19 @@ fn speak_proxy(req: Request, remote_addr: SocketAddr, endpoint: &'static s speaker: request.speaker, } }).and_then(move |request_details: RequestDetails| { + let proxy_host = match router.get_speaker_host(&request_details.speaker) { + Some(host) => host, + None => { + let response = Response::builder() + .body(Body::from("error")) + .unwrap(); + let result : BoxFut = Box::new(future::ok(response)); + return result; + }, + }; + + info!("Routing {} for {} to {}", endpoint, &request_details.speaker, &proxy_host); + let mut request_builder = Request::builder(); request_builder.method(Method::POST) .uri(endpoint); @@ -128,7 +157,7 @@ fn speak_proxy(req: Request, remote_addr: SocketAddr, endpoint: &'static s .body(Body::from(request_details.request_bytes)) .unwrap(); - hyper_reverse_proxy::call(remote_addr.ip(), "http://127.0.0.1:12345", new_req) + hyper_reverse_proxy::call(remote_addr.ip(), &proxy_host, new_req) })) } @@ -145,12 +174,10 @@ fn main() -> AnyhowResult<()> { env_logger::init(); - let route_one = get_env_string(ENV_ROUTE_ONE, ROUTE_ONE_DEFAULT); - let route_two = get_env_string(ENV_ROUTE_TWO, ROUTE_TWO_DEFAULT); + let default_route = get_env_string(ENV_DEFAULT_ROUTE, DEFAULT_DEFAULT_ROUTE); let service_port = get_env_num::(ENV_SERVICE_PORT, DEFAULT_SERVICE_PORT)?; - info!("Route one: {}", route_one); - info!("Route two: {}", route_two); + info!("Default route: {}", default_route); let proxy_configs_file = get_env_string(ENV_PROXY_CONFIG_FILE, DEFAULT_PROXY_CONFIG_FILE); info!("Proxy config file: {}", proxy_configs_file); @@ -158,33 +185,51 @@ fn main() -> AnyhowResult<()> { let proxy_configs = ProxyConfigs::load_from_file(&proxy_configs_file)?; info!("Proxy configs: {:?}", proxy_configs); - // This is our socket address... - let addr = ([127, 0, 0, 1], service_port).into(); + let mut route_map = HashMap::new(); + + for backend in proxy_configs.backends.iter() { + let host = format!("http://{}:{}", backend.ip, backend.port); + route_map.insert(backend.voice.clone(), host); + } + + let router = Arc::new(Router { + routes: route_map, + }); + + //let routes = route_map.clone(); // A `Service` is needed for every connection. - let make_svc = make_service_fn(move |socket: &AddrStream| { + let make_service = make_service_fn(move |socket: &AddrStream| { let remote_addr = socket.remote_addr(); info!("Got a request from: {:?}", remote_addr); - let route_a = route_one.clone(); + let router2 = router.clone(); + let default_route2 = default_route.clone(); service_fn(move |req: Request| { + let router3 = router2.clone(); + match (req.method(), req.uri().path()) { (&Method::POST, "/speak") => - speak_proxy(req, remote_addr.clone(), "/speak"), + speak_proxy(req, remote_addr.clone(), router3, "/speak"), (&Method::POST, "/speak_spectrogram") => - speak_proxy(req, remote_addr.clone(), "/speak_spectrogram"), - _ => - hyper_reverse_proxy::call(remote_addr.ip(), &route_a, req), + speak_proxy(req, remote_addr.clone(), router3, "/speak_spectrogram"), + _ => { + let forward = router3.get_random_host() + .unwrap_or(default_route2.clone()); + hyper_reverse_proxy::call(remote_addr.ip(), &forward, req) + } } }) }); - let server = Server::bind(&addr) - .serve(make_svc) - .map_err(|e| eprintln!("server error: {}", e)); + let addr = ([127, 0, 0, 1], service_port).into(); - println!("Running server on {:?}", addr); + info!("Running server on {:?}", addr); + + let server = Server::bind(&addr) + .serve(make_service) + .map_err(|e| eprintln!("server error: {}", e)); // Run this server for... forever! hyper::rt::run(server);