diff --git a/src/main.rs b/src/main.rs index 8d3a903..09f71bb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,12 +5,13 @@ #[macro_use] extern crate serde_derive; pub mod newrelic_logger; +pub mod request; pub mod speak_proxy; use anyhow::Result as AnyhowResult; use crate::newrelic_logger::{NewRelicLogger, MaybeNewRelicTransaction}; use futures::future::{self, Future}; -use crate::speak_proxy::{speak_proxy, speak_proxy_with_retry}; +use crate::speak_proxy::speak_proxy_with_retry; use hyper::{Method, StatusCode}; use hyper::header::HeaderValue; use hyper::http::header::HeaderName; diff --git a/src/request.rs b/src/request.rs new file mode 100644 index 0000000..7e3d628 --- /dev/null +++ b/src/request.rs @@ -0,0 +1,65 @@ +use anyhow::Result as AnyhowResult; +use futures::TryStreamExt; +use hyper::http::HeaderValue; +use hyper::{Body, Request, HeaderMap}; + +/// NB: This much match the shape of SpeakRequest in the 'voder/tts_service' code. +/// This is used for both /speak and /speak_spectrogram requests. +#[derive(Deserialize, Debug)] +pub struct ExternalSpeakRequest { + /// Slug for the speaker + pub speaker: String, + + /// Raw text to be spoken + pub text: String, +} + +/// The full request +pub struct ExternalHttpRequest { + pub headers: HeaderMap, + pub speak_request: ExternalSpeakRequest, +} + +#[derive(Serialize, Debug)] +pub struct InternalSpeakRequest { + /// Slug for the speaker + pub speaker: String, + + /// Raw text to be spoken + pub text: String, + + /// Retry attempt. + /// The first request is "0", and all subsequent requests are incremented. + /// Subsequent requests do not trigger the rate limiter. + pub retry_attempt_number: u8, + + /// Admin/debug requests get to skip the rate limiter. + pub skip_rate_limiter: bool, +} + + +impl ExternalHttpRequest { + + pub async fn decode_request(request: Request) -> AnyhowResult { + let headers = request.headers().clone(); + let (_parts, body) = request.into_parts(); + + let body_bytes = body_to_bytes(body).await?; + let speak_request = serde_json::from_slice::(&body_bytes)?; + + Ok(ExternalHttpRequest { + headers: headers, + speak_request: speak_request, + }) + } +} + +async fn body_to_bytes(body: Body) -> AnyhowResult> { + // Also: let body_bytes = hyper::body::to_bytes(body) + let bytes = body.try_fold(Vec::new(), |mut data, chunk| async move { + data.extend_from_slice(&chunk); + Ok(data) + }).await?; + + Ok(bytes) +} diff --git a/src/speak_proxy.rs b/src/speak_proxy.rs index 9e8c7d5..aa77746 100644 --- a/src/speak_proxy.rs +++ b/src/speak_proxy.rs @@ -7,7 +7,7 @@ use hyper::header::HeaderValue; use hyper::http::header::HeaderName; use hyper::server::conn::AddrStream; use hyper::service::{service_fn, make_service_fn}; -use hyper::{Body, Request, Response, Server}; +use hyper::{Body, Request, Response, Server, HeaderMap}; use hyper::{Method, StatusCode}; use hyper_reverse_proxy::ProxyError; use rand::seq::IteratorRandom; @@ -19,58 +19,61 @@ use std::fmt::Display; use std::net::{SocketAddr, IpAddr}; use std::str::FromStr; use std::sync::Arc; - -/// NB: This much match the shape of SpeakRequest in the 'voder/tts_service' code. -/// This is used for both /speak and /speak_spectrogram requests. -#[derive(Serialize, Deserialize, Debug)] -pub struct SpeakRequest { - /// Slug for the speaker - speaker: String, - /// Raw text to be spoken - text: String, -} - -struct RequestDetails { - pub request_bytes: Vec, - pub speaker: String, -} +use crate::request::{ExternalSpeakRequest, ExternalHttpRequest, InternalSpeakRequest}; pub async fn speak_proxy_with_retry( ip_addr: IpAddr, - req: Request, + request: Request, router: Arc, _transaction: MaybeNewRelicTransaction, endpoint: &'static str) -> Result, ResponseError> { let mut attempt : u8 = 0; - let headers = req.headers().clone(); + let external_http_request = ExternalHttpRequest::decode_request(request).await.unwrap(); - let (_parts, body) = req.into_parts(); + // This is the IP of the Digital Ocean load balancer. + let gateway_ip = HeaderValue::from_str(&ip_addr.to_string()) + .ok() + .unwrap_or(HeaderValue::from_static("")); - let body_bytes = body_to_bytes(body).await.unwrap(); + // This is the IP of the upstream client browser. + let upstream_client_ip = external_http_request.headers.get(HeaderName::from_static("x-forwarded-for")) + .map(|ip| ip.to_str().unwrap_or("")) + .unwrap_or(""); - let body_string = String::from_utf8(body_bytes) - .unwrap(); + let upstream_client_ip = HeaderValue::from_str(upstream_client_ip) + .ok() + .unwrap_or(HeaderValue::from_static("")); - println!("Body: {:?}", body_string); + let mut headers_to_proxy = external_http_request.headers.clone(); + + headers_to_proxy.insert(HeaderName::from_static("forwarded"), upstream_client_ip.clone()); + headers_to_proxy.insert(HeaderName::from_static("x-forwarded-for"), upstream_client_ip.clone()); + + // Unfortunately it looks like this middleware trounces the standard headers, so + // here we put it somewhere it won't be overwritten. + headers_to_proxy.insert(HeaderName::from_static("x-voder-proxy-for"), upstream_client_ip.clone()); + headers_to_proxy.insert(HeaderName::from_static("x-gateway-ip"), gateway_ip); + + // The proxy host is a Kubernetes load balancer. We're not doing the round robin magic here, + // so it makes sense to reuse the same value for each retry. + let load_balancer_hostname = router.get_speaker_host(&external_http_request.speak_request.speaker); loop { - info!("Attempt #{}", attempt); + let internal_speak_request = InternalSpeakRequest { + speaker: external_http_request.speak_request.speaker.clone(), + text: external_http_request.speak_request.text.clone(), + retry_attempt_number: attempt, + skip_rate_limiter: false, + }; - let mut request_builder = Request::builder(); - - for (k, v) in headers.iter() { - // NB: Builder pattern ownership - request_builder = request_builder.header(k.clone(), v.clone()); - } - - let request = request_builder - .body(Body::from(body_string.clone())) - .unwrap(); - - let transaction = MaybeNewRelicTransaction::from_optional(None); - let result = speak_proxy(ip_addr, request, router.clone(), transaction, endpoint).await; + let result = speak_proxy( + ip_addr, + internal_speak_request, + headers_to_proxy.clone(), + &load_balancer_hostname, + endpoint).await; if result.is_ok() { return result; @@ -89,53 +92,14 @@ pub struct ResponseError { pub reason: String, } -pub async fn speak_proxy( +async fn speak_proxy( ip_addr: IpAddr, - req: Request, - router: Arc, - transaction: MaybeNewRelicTransaction, + internal_speak_request: InternalSpeakRequest, + headers: HeaderMap, + load_balancer_hostname: &str, endpoint: &'static str) -> Result, ResponseError> { - let mut headers = req.headers().clone(); - - // This is the IP of the Digital Ocean load balancer. - let gateway_ip = HeaderValue::from_str(&ip_addr.to_string()) - .ok() - .unwrap_or(HeaderValue::from_static("")); - - // This is the IP of the upstream client browser. - let upstream_client_ip = req.headers().get(HeaderName::from_static("x-forwarded-for")) - .map(|ip| ip.to_str().unwrap_or("")) - .unwrap_or(""); - - let upstream_client_ip = HeaderValue::from_str(upstream_client_ip) - .ok() - .unwrap_or(HeaderValue::from_static("")); - - headers.insert(HeaderName::from_static("forwarded"), upstream_client_ip.clone()); - headers.insert(HeaderName::from_static("x-forwarded-for"), upstream_client_ip.clone()); - // Unfortunately it looks like this middleware trounces the standard headers, so - // here we put it somewhere it won't be overwritten. - headers.insert(HeaderName::from_static("x-voder-proxy-for"), upstream_client_ip.clone()); - headers.insert(HeaderName::from_static("x-gateway-ip"), gateway_ip); - - let (_parts, body) = req.into_parts(); - - let request_bytes = body_to_bytes(body).await.unwrap(); - - let request = serde_json::from_slice::(&request_bytes) - .unwrap(); - - let request_details = RequestDetails { - request_bytes: request_bytes.clone(), - speaker: request.speaker, - }; - - let proxy_host = router.get_speaker_host(&request_details.speaker); - - info!("Routing {} for {} to {}", endpoint, &request_details.speaker, &proxy_host); - - transaction.add_attribute("speaker", &request_details.speaker); + info!("Routing {} for {} to {}", endpoint, &internal_speak_request.speaker, &load_balancer_hostname); let mut request_builder = Request::builder() .method(Method::POST) @@ -146,26 +110,17 @@ pub async fn speak_proxy( request_builder = request_builder.header(k, v); } - let new_req = request_builder - .body(Body::from(request_details.request_bytes)) + let json_speak_request = serde_json::to_string(&internal_speak_request).unwrap(); + + let request = request_builder + .body(Body::from(json_speak_request)) .unwrap(); //Box::new(future::ok(Vec::new())) - match hyper_reverse_proxy::call(ip_addr, &proxy_host, new_req).await { + match hyper_reverse_proxy::call(ip_addr, &load_balancer_hostname, request).await { Ok(result) => Ok(result), Err(_result) => Err(ResponseError { reason: "Uknown failure".to_string() }), } } - -async fn body_to_bytes(body: Body) -> AnyhowResult> { - // Also: let body_bytes = hyper::body::to_bytes(body) - let bytes = body.try_fold(Vec::new(), |mut data, chunk| async move { - data.extend_from_slice(&chunk); - Ok(data) - }) - .await?; - - Ok(bytes) -}