Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 38 additions & 4 deletions src/web/helper.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -218,21 +218,34 @@ 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<Cookie<'a>>,
}

// No inline needed... generic are inherently inlined
#[allow(clippy::too_many_arguments)]
pub async fn forward<T>(
req: HttpRequest,
mut payload: web::Payload,
method: Method,
peer_addr: Option<PeerAddr>,
client: web::Data<reqwest::Client>,
new_url: T,
updated_cookie: Option<Cookie<'_>>,
inference: bool,
config: Config<'_>,
) -> Result<HttpResponse>
where
T: AsRef<str> + 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 {
Expand Down Expand Up @@ -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| {
Expand Down
2 changes: 1 addition & 1 deletion src/web/route/health/inference_services.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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())
})??;

Expand Down
4 changes: 2 additions & 2 deletions src/web/route/mcp/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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/";
Expand Down Expand Up @@ -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");
Expand Down
12 changes: 10 additions & 2 deletions src/web/route/notebook/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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},
},
};

Expand Down Expand Up @@ -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
}
Expand Down
28 changes: 25 additions & 3 deletions src/web/route/openwebui/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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))]
Expand All @@ -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]
Expand Down
19 changes: 17 additions & 2 deletions src/web/route/proxy/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down
13 changes: 10 additions & 3 deletions src/web/route/resource/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
}
Expand Down