diff --git a/src/web/helper.rs b/src/web/helper.rs index 1a7bbfd6..b97f68e9 100644 --- a/src/web/helper.rs +++ b/src/web/helper.rs @@ -1,8 +1,8 @@ use actix_web::web; use base64::{Engine, prelude::BASE64_STANDARD}; -use mongodb::bson::{Bson, to_bson}; #[cfg(feature = "observe")] use mongodb::bson; +use mongodb::bson::{Bson, to_bson}; use rand::{Rng, rng}; use serde::Deserialize; use tera::Context; @@ -218,8 +218,18 @@ pub mod forwarding { use crate::errors::{BridgeError, Result}; + /// Configuration that is passed to the forward function takes the following parameters. + /// inference filters out Auth specific headers to the proxy forward + /// pack_cookies packs cookies to be compatible for http protocol older than 2 + /// updated_cookie will forward back to client any updated cookies + #[derive(Default)] + pub struct Config<'a> { + pub inference: bool, + pub pack_cookies: bool, + pub updated_cookie: Option>, + } + // No inline needed... generic are inherently inlined - #[allow(clippy::too_many_arguments)] pub async fn forward( req: HttpRequest, mut payload: web::Payload, @@ -227,12 +237,15 @@ pub mod forwarding { peer_addr: Option, client: web::Data, new_url: T, - updated_cookie: Option>, - inference: bool, + config: Config<'_>, ) -> Result where T: AsRef + Send + Sync, { + let inference = config.inference; + let pack_cookies = config.pack_cookies; + let updated_cookie = config.updated_cookie; + let (tx, rx) = mpsc::channel(128); actix_web::rt::spawn(async move { @@ -292,6 +305,27 @@ pub mod forwarding { ); } + // find all the cookie header and pack them delimited by ";" + if pack_cookies { + let mut cookies = String::new(); + for (header_name, header_value) in req.headers().iter() { + if header_name.as_str().to_lowercase() == "cookie" { + if let Ok(value) = header_value.to_str() { + if !cookies.is_empty() { + cookies.push(';'); + } + cookies.push_str(value); + } + } + } + if !cookies.is_empty() { + headers.insert( + ReqwestHeaderName::from_static("cookie"), + ReqwestHeaderValue::from_str(&cookies).unwrap(), + ); + } + } + let forwarded_req = forwarded_req.headers(headers); let res = forwarded_req.send().await.map_err(|e| { diff --git a/src/web/route/health/inference_services.rs b/src/web/route/health/inference_services.rs index 1e291908..ee482ec5 100644 --- a/src/web/route/health/inference_services.rs +++ b/src/web/route/health/inference_services.rs @@ -67,7 +67,7 @@ impl<'a> InferenceServicesHealth<'a> { let now = Instant::now(); let fut = client.get(url.as_str()).send(); - let response = timeout(Duration::from_secs(1), fut).await.map_err(|_| { + let response = timeout(Duration::from_secs(10), fut).await.map_err(|_| { BridgeError::GeneralError("Call to inference service timed out".to_string()) })??; diff --git a/src/web/route/mcp/mod.rs b/src/web/route/mcp/mod.rs index e2afa793..b5eb97fc 100644 --- a/src/web/route/mcp/mod.rs +++ b/src/web/route/mcp/mod.rs @@ -10,7 +10,7 @@ use crate::{ auth::jwt::validate_token, config::CONFIG, errors::{BridgeError, Result}, - web::{helper, services::CATALOG}, + web::{helper::{self, forwarding::Config}, services::CATALOG}, }; const MCP_PREFIX: &str = "/mcp/"; @@ -68,7 +68,7 @@ async fn forward( } new_url.set_query(req.uri().query()); - helper::forwarding::forward(req, payload, method, peer_addr, client, new_url, None, true) + helper::forwarding::forward(req, payload, method, peer_addr, client, new_url, Config{inference: true, ..Default::default()}) .await } else { warn!("MCP service not found in url request"); diff --git a/src/web/route/notebook/mod.rs b/src/web/route/notebook/mod.rs index e6fae9d7..722f9900 100644 --- a/src/web/route/notebook/mod.rs +++ b/src/web/route/notebook/mod.rs @@ -37,7 +37,7 @@ use crate::{ kube::{KubeAPI, NOTEBOOK_NAMESPACE, Notebook, NotebookSpec, PVCSpec}, web::{ bridge_middleware::{CookieCheck, Htmx, NotebookCookieCheck}, - helper::{self, bson}, + helper::{self, bson, forwarding}, }, }; @@ -672,7 +672,15 @@ async fn notebook_forward( new_url.set_query(req.uri().query()); helper::forwarding::forward( - req, payload, method, peer_addr, client, new_url, None, false, + req, + payload, + method, + peer_addr, + client, + new_url, + forwarding::Config { + ..Default::default() + }, ) .await } diff --git a/src/web/route/openwebui/mod.rs b/src/web/route/openwebui/mod.rs index e65d45fa..53cfeb2a 100644 --- a/src/web/route/openwebui/mod.rs +++ b/src/web/route/openwebui/mod.rs @@ -14,7 +14,7 @@ use crate::{ config::CONFIG, db::models::OWUICookie, errors::{BridgeError, Result}, - web::helper, + web::helper::{self, forwarding}, }; const OWUI_PORT: &str = "8080"; @@ -90,7 +90,18 @@ async fn openwebui_forward( url.set_query(req.uri().query()); - helper::forwarding::forward(req, payload, method, peer_addr, client, url, None, false).await + helper::forwarding::forward( + req, + payload, + method, + peer_addr, + client, + url, + forwarding::Config { + ..Default::default() + }, + ) + .await } #[instrument(skip(payload))] @@ -116,7 +127,18 @@ async fn moleviewer_forward( url.set_path(path); url.set_query(req.uri().query()); - helper::forwarding::forward(req, payload, method, peer_addr, client, url, None, false).await + helper::forwarding::forward( + req, + payload, + method, + peer_addr, + client, + url, + forwarding::Config { + ..Default::default() + }, + ) + .await } #[inline] diff --git a/src/web/route/proxy/mod.rs b/src/web/route/proxy/mod.rs index 7478d8cb..cf35cc84 100644 --- a/src/web/route/proxy/mod.rs +++ b/src/web/route/proxy/mod.rs @@ -4,7 +4,10 @@ use tracing::{error, instrument, warn}; use crate::{ errors::{BridgeError, Result}, - web::{bridge_middleware::validator, helper}, + web::{ + bridge_middleware::validator, + helper::{self, forwarding::Config}, + }, }; use self::services::CATALOG; @@ -47,7 +50,19 @@ async fn forward( new_url.set_path(path); new_url.set_query(req.uri().query()); - helper::forwarding::forward(req, payload, method, peer_addr, client, new_url, None, true).await + helper::forwarding::forward( + req, + payload, + method, + peer_addr, + client, + new_url, + Config { + inference: true, + ..Default::default() + }, + ) + .await } pub fn config_proxy(cfg: &mut web::ServiceConfig) { diff --git a/src/web/route/resource/mod.rs b/src/web/route/resource/mod.rs index 4bde19bf..9bdf9216 100644 --- a/src/web/route/resource/mod.rs +++ b/src/web/route/resource/mod.rs @@ -19,7 +19,11 @@ use crate::{ mongo::DB, }, errors::{BridgeError, Result}, - web::{bridge_middleware::ResourceCookieCheck, helper, services::CATALOG}, + web::{ + bridge_middleware::ResourceCookieCheck, + helper::{self, forwarding::Config}, + services::CATALOG, + }, }; static TOKEN_LIFETIME: usize = 60 * 60 * 24; // 24 hours @@ -88,8 +92,11 @@ async fn resource_http( peer_addr, client, new_url, - updated_cookie, - false, + Config { + updated_cookie, + pack_cookies: true, + ..Default::default() + }, ) .await }