use std::cell::RefCell;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use tokio::io::{AsyncReadExt, AsyncWriteExt, DuplexStream, duplex};
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use super::*;
use crate::token_policy::TokenPolicy;
const TOKEN: &str = "sk-zzq-the-credential";
const TEST_PEER: &str = "3ca82708b995";
const OK_RESPONSE: &[u8] =
b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 2\r\n\r\n{}";
struct CountingBackend {
connects: Arc<AtomicUsize>,
response: Vec<u8>,
received: Arc<Mutex<Vec<JoinHandle<Vec<u8>>>>>,
}
impl CountingBackend {
fn new(response: &[u8]) -> Self {
Self {
connects: Arc::new(AtomicUsize::new(0)),
response: response.to_vec(),
received: Arc::new(Mutex::new(Vec::new())),
}
}
fn connects(&self) -> usize {
self.connects.load(Ordering::SeqCst)
}
async fn received(&self) -> Vec<u8> {
let taken: Vec<_> = self.received.lock().await.drain(..).collect();
let mut out = Vec::new();
for handle in taken {
out.extend_from_slice(&handle.await.expect("backend task"));
}
out
}
}
impl Backend for CountingBackend {
type Stream = DuplexStream;
fn authority(&self) -> &'static str {
"127.0.0.1:11434"
}
async fn connect(&self) -> std::io::Result<DuplexStream> {
self.connects.fetch_add(1, Ordering::SeqCst);
let (mine, mut theirs) = duplex(64 * 1024);
let response = self.response.clone();
let handle = tokio::spawn(async move {
let _ = theirs.write_all(&response).await;
let _ = theirs.flush().await;
let mut seen = Vec::new();
let _ = theirs.read_to_end(&mut seen).await;
seen
});
self.received.lock().await.push(handle);
Ok(mine)
}
}
struct Fixed(Mutex<Option<DuplexStream>>);
impl Backend for Fixed {
type Stream = DuplexStream;
fn authority(&self) -> &'static str {
"127.0.0.1:11434"
}
async fn connect(&self) -> std::io::Result<DuplexStream> {
Ok(self.0.lock().await.take().expect("connected once"))
}
}
async fn exchange(
request: &[u8],
policy: &TokenPolicy,
backend: &CountingBackend,
) -> (Outcome, Vec<u8>) {
let (mut client, mut edge) = duplex(64 * 1024);
client.write_all(request).await.unwrap();
client.shutdown().await.unwrap();
let (credential, _) = Credential::new(policy).expect("a usable policy");
let outcome = tokio::time::timeout(
std::time::Duration::from_secs(5),
serve_exchange(&mut edge, &credential, backend, TEST_PEER),
)
.await
.expect("the exchange must not hang")
.expect("no transport failure");
drop(edge);
let mut seen = Vec::new();
client.read_to_end(&mut seen).await.unwrap();
(outcome, seen)
}
fn get(auth: Option<&str>) -> Vec<u8> {
let mut req = b"GET /v1/models HTTP/1.1\r\nHost: 127.0.0.1:8080\r\n".to_vec();
if let Some(value) = auth {
req.extend_from_slice(format!("Authorization: {value}\r\n").as_bytes());
}
req.extend_from_slice(b"\r\n");
req
}
fn supplied() -> TokenPolicy {
TokenPolicy::Supplied(TOKEN.to_owned())
}
#[tokio::test]
async fn an_authorized_request_reaches_the_backend_and_the_response_comes_back() {
let backend = CountingBackend::new(OK_RESPONSE);
let (outcome, seen) = exchange(
&get(Some(&format!("Bearer {TOKEN}"))),
&supplied(),
&backend,
)
.await;
assert_eq!(outcome, Outcome::Forwarded);
assert_eq!(backend.connects(), 1);
assert!(
String::from_utf8_lossy(&seen).starts_with("HTTP/1.1 200 OK"),
"the backend's response reaches the client: {}",
String::from_utf8_lossy(&seen)
);
}
#[tokio::test]
async fn serving_open_forwards_a_request_with_no_credential() {
let backend = CountingBackend::new(OK_RESPONSE);
let (outcome, _) = exchange(&get(None), &TokenPolicy::InsecureNoAuth, &backend).await;
assert_eq!(outcome, Outcome::Forwarded);
assert_eq!(backend.connects(), 1);
}
#[tokio::test]
async fn a_rejected_request_opens_no_backend_connection_at_all() {
let backend = CountingBackend::new(OK_RESPONSE);
let wrong = [
None,
Some("Bearer wrong-token"),
Some(TOKEN),
Some(&format!("Bearer {TOKEN}x")),
Some(&format!("Basic {TOKEN}")),
];
for (i, auth) in wrong.iter().enumerate() {
let (outcome, seen) = exchange(&get(auth.as_deref()), &supplied(), &backend).await;
assert_eq!(outcome, Outcome::Unauthorized, "case {i}");
assert!(
String::from_utf8_lossy(&seen).starts_with("HTTP/1.1 401"),
"case {i} must be refused"
);
}
assert_eq!(
backend.connects(),
0,
"after {} unauthorized requests the backend was never contacted",
wrong.len()
);
assert!(backend.received().await.is_empty(), "and sent nothing");
}
#[tokio::test]
async fn the_401_is_synthesized_locally_and_advertises_the_scheme() {
let backend = CountingBackend::new(OK_RESPONSE);
let (_, seen) = exchange(&get(None), &supplied(), &backend).await;
let text = String::from_utf8_lossy(&seen);
assert!(text.starts_with("HTTP/1.1 401 Unauthorized"));
assert!(text.contains("WWW-Authenticate: Bearer"));
assert!(text.contains("invalid_api_key"));
assert_eq!(backend.connects(), 0);
}
#[tokio::test]
async fn an_ambiguously_framed_request_is_refused_without_a_backend_connection() {
let backend = CountingBackend::new(OK_RESPONSE);
let smuggle = b"POST /v1/chat/completions HTTP/1.1\r\nHost: x\r\nContent-Length: 6\r\nTransfer-Encoding: chunked\r\n\r\n0\r\n\r\n";
let (outcome, seen) = exchange(smuggle, &supplied(), &backend).await;
assert_eq!(outcome, Outcome::BadRequest);
assert!(String::from_utf8_lossy(&seen).starts_with("HTTP/1.1 400"));
assert_eq!(backend.connects(), 0, "smuggling never reaches the backend");
}
#[tokio::test]
async fn framing_is_refused_before_the_credential_is_consulted() {
let backend = CountingBackend::new(OK_RESPONSE);
let mut smuggle = b"POST /v1/chat/completions HTTP/1.1\r\nHost: x\r\n".to_vec();
smuggle.extend_from_slice(format!("Authorization: Bearer {TOKEN}\r\n").as_bytes());
smuggle.extend_from_slice(b"Content-Length: 6\r\nTransfer-Encoding: chunked\r\n\r\n0\r\n\r\n");
let (outcome, _) = exchange(&smuggle, &supplied(), &backend).await;
assert_eq!(outcome, Outcome::BadRequest);
assert_eq!(backend.connects(), 0);
}
#[tokio::test]
async fn a_head_that_is_not_http_is_refused_without_a_backend_connection() {
let backend = CountingBackend::new(OK_RESPONSE);
let (outcome, seen) = exchange(b"this is not http\r\n\r\n", &supplied(), &backend).await;
assert_eq!(outcome, Outcome::BadRequest);
assert!(String::from_utf8_lossy(&seen).starts_with("HTTP/1.1 400"));
assert_eq!(backend.connects(), 0);
}
#[tokio::test]
async fn the_backend_receives_the_message_headers_and_not_the_connection_ones() {
let backend = CountingBackend::new(OK_RESPONSE);
let mut req = b"GET /v1/models HTTP/1.1\r\nHost: 127.0.0.1:8080\r\n".to_vec();
req.extend_from_slice(format!("Authorization: Bearer {TOKEN}\r\n").as_bytes());
req.extend_from_slice(b"Connection: keep-alive, X-Hop\r\nX-Hop: 1\r\n");
req.extend_from_slice(b"X-Forwarded-For: 203.0.113.1\r\nAccept: */*\r\n\r\n");
let (outcome, _) = exchange(&req, &supplied(), &backend).await;
assert_eq!(outcome, Outcome::Forwarded);
let sent = String::from_utf8(backend.received().await).expect("ascii");
let lower = sent.to_ascii_lowercase();
assert!(
sent.contains("Host: 127.0.0.1:11434"),
"the backend's own authority: {sent}"
);
assert!(!sent.contains("127.0.0.1:8080"), "not the client's: {sent}");
assert!(
lower.contains("authorization: bearer"),
"forwarded for the second check"
);
assert!(
!lower.contains("x-hop"),
"a nominated hop-by-hop header: {sent}"
);
assert!(
!lower.contains("connection: keep-alive"),
"the client's connection header must not survive: {sent}"
);
assert_eq!(
lower.matches("connection:").count(),
1,
"exactly one, and it is the edge's: {sent}"
);
assert!(
lower.contains("connection: close"),
"the backend is told this connection carries one exchange: {sent}"
);
assert!(
!lower.contains("x-forwarded-for"),
"a forwarding chain: {sent}"
);
assert!(lower.contains("accept: */*"), "a message header survives");
}
#[tokio::test]
async fn the_backend_sees_the_edges_tunnel_markers_and_not_the_clients() {
let backend = CountingBackend::new(OK_RESPONSE);
let mut req = b"GET /v1/models HTTP/1.1\r\nHost: 127.0.0.1:8080\r\n".to_vec();
req.extend_from_slice(format!("Authorization: Bearer {TOKEN}\r\n").as_bytes());
req.extend_from_slice(b"Via: 1.1 somebody-else\r\nVIA: 1.0 another\r\n");
req.extend_from_slice(b"X-Modelpipe-Peer: 000000000000\r\n\r\n");
let (outcome, _) = exchange(&req, &supplied(), &backend).await;
assert_eq!(outcome, Outcome::Forwarded);
let sent = String::from_utf8(backend.received().await).expect("ascii");
let lower = sent.to_ascii_lowercase();
assert_eq!(
lower.matches("\r\nvia:").count(),
1,
"exactly one Via: {sent}"
);
assert!(
sent.contains("Via: 1.1 modelpipe"),
"and it is the edge's: {sent}"
);
assert!(
!lower.contains("somebody-else") && !lower.contains("another"),
"{sent}"
);
assert_eq!(
lower.matches("\r\nx-modelpipe-peer:").count(),
1,
"exactly one peer marker: {sent}"
);
assert!(
sent.contains(&format!("X-Modelpipe-Peer: {TEST_PEER}")),
"and it names the peer the listener saw: {sent}"
);
assert!(
!sent.contains("000000000000"),
"not the one the client claimed: {sent}"
);
}
#[tokio::test]
async fn a_request_body_reaches_the_backend_byte_for_byte() {
let backend = CountingBackend::new(OK_RESPONSE);
let payload = r#"{"model":"llama","messages":[{"role":"user","content":"hi"}]}"#;
let mut req = b"POST /v1/chat/completions HTTP/1.1\r\nHost: x\r\n".to_vec();
req.extend_from_slice(format!("Authorization: Bearer {TOKEN}\r\n").as_bytes());
req.extend_from_slice(format!("Content-Length: {}\r\n\r\n", payload.len()).as_bytes());
req.extend_from_slice(payload.as_bytes());
let (outcome, _) = exchange(&req, &supplied(), &backend).await;
assert_eq!(outcome, Outcome::Forwarded);
assert!(
String::from_utf8(backend.received().await)
.expect("ascii")
.ends_with(payload),
"the body must arrive unaltered"
);
}
#[tokio::test]
async fn a_streaming_response_reaches_the_client_frame_by_frame() {
let backend = CountingBackend::new(OK_RESPONSE);
let (mut client, mut edge) = duplex(64 * 1024);
client
.write_all(&get(Some(&format!("Bearer {TOKEN}"))))
.await
.unwrap();
let (mine, mut theirs) = duplex(64 * 1024);
let released = Arc::new(tokio::sync::Notify::new());
let wait = released.clone();
tokio::spawn(async move {
theirs
.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\r\n")
.await
.unwrap();
theirs.write_all(b"data: first\n\n").await.unwrap();
theirs.flush().await.unwrap();
wait.notified().await;
theirs.write_all(b"data: [DONE]\n\n").await.unwrap();
});
let fixed = Fixed(Mutex::new(Some(mine)));
let (credential, _) = Credential::new(&supplied()).expect("a usable token");
let pump = tokio::spawn(async move {
serve_exchange(&mut edge, &credential, &fixed, TEST_PEER)
.await
.unwrap()
});
let mut text = String::new();
tokio::time::timeout(std::time::Duration::from_secs(5), async {
let mut buf = [0u8; 512];
loop {
let n = client.read(&mut buf).await.expect("read");
assert!(n > 0, "the stream ended early: {text}");
text.push_str(&String::from_utf8_lossy(&buf[..n]));
if text.contains("data: first") {
return;
}
}
})
.await
.unwrap_or_else(|_| panic!("the first frame must arrive before the backend finishes: {text}"));
assert!(text.contains("200 OK"), "{text}");
released.notify_one();
assert_eq!(pump.await.unwrap(), Outcome::Forwarded);
let _ = backend.connects();
}
struct KeepAlive(Mutex<Option<DuplexStream>>);
impl KeepAlive {
fn new(response: &'static str) -> Self {
let (mine, mut theirs) = duplex(64 * 1024);
tokio::spawn(async move {
let mut sink = Vec::new();
let _ = theirs.write_all(response.as_bytes()).await;
let _ = theirs.flush().await;
let _ = theirs.read_to_end(&mut sink).await;
});
Self(Mutex::new(Some(mine)))
}
}
impl Backend for KeepAlive {
type Stream = DuplexStream;
fn authority(&self) -> &'static str {
"127.0.0.1:11434"
}
async fn connect(&self) -> std::io::Result<DuplexStream> {
Ok(self.0.lock().await.take().expect("connected once"))
}
}
async fn against_keepalive(request: &[u8], response: &'static str) -> (Outcome, String) {
let backend = KeepAlive::new(response);
let (mut client, mut edge) = duplex(64 * 1024);
client.write_all(request).await.unwrap();
let (credential, _) = Credential::new(&supplied()).expect("a usable token");
let outcome = tokio::time::timeout(
std::time::Duration::from_secs(5),
serve_exchange(&mut edge, &credential, &backend, TEST_PEER),
)
.await
.expect("a keep-alive backend must not hang the exchange")
.expect("no transport failure");
drop(edge);
let mut seen = Vec::new();
client.read_to_end(&mut seen).await.unwrap();
(outcome, String::from_utf8_lossy(&seen).into_owned())
}
fn authed(method: &str) -> Vec<u8> {
format!("{method} /v1/models HTTP/1.1\r\nHost: x\r\nAuthorization: Bearer {TOKEN}\r\n\r\n")
.into_bytes()
}
#[tokio::test]
async fn a_bodyless_status_is_framed_by_its_status_and_not_by_its_headers() {
for response in [
"HTTP/1.1 204 No Content\r\n\r\n",
"HTTP/1.1 304 Not Modified\r\nContent-Length: 42\r\n\r\n",
] {
let (outcome, seen) = against_keepalive(&authed("GET"), response).await;
assert_eq!(outcome, Outcome::Forwarded, "{response:?}");
assert!(
seen.starts_with("HTTP/1.1 3") || seen.starts_with("HTTP/1.1 2"),
"{seen}"
);
}
}
#[tokio::test]
async fn a_head_response_is_not_a_body_to_wait_for() {
let (outcome, seen) = against_keepalive(
&authed("HEAD"),
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 4096\r\n\r\n",
)
.await;
assert_eq!(outcome, Outcome::Forwarded);
assert!(seen.starts_with("HTTP/1.1 200"), "{seen}");
assert!(
seen.to_ascii_lowercase().contains("content-length: 4096"),
"the declared length still describes what a GET would return: {seen}"
);
}
#[tokio::test]
async fn an_interim_response_is_skipped_and_the_real_one_forwarded() {
let (outcome, seen) = against_keepalive(
&authed("GET"),
"HTTP/1.1 100 Continue\r\n\r\nHTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\n{}",
)
.await;
assert_eq!(outcome, Outcome::Forwarded);
assert!(
seen.starts_with("HTTP/1.1 200"),
"the interim head must not be the answer: {seen}"
);
assert!(seen.ends_with("{}"), "and the real body arrives: {seen}");
}
#[tokio::test]
async fn an_ambiguously_framed_backend_response_is_refused_rather_than_resolved() {
let (outcome, seen) = against_keepalive(
&authed("GET"),
"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nTransfer-Encoding: chunked\r\n\r\n{}",
)
.await;
assert_eq!(outcome, Outcome::BadGateway);
assert!(
seen.starts_with("HTTP/1.1 502"),
"the client did nothing wrong, and the backend's answer is unreadable: {seen}"
);
}
#[tokio::test]
async fn a_backend_that_refuses_the_connection_is_reported_as_a_gateway_failure() {
struct Refusing;
impl Backend for Refusing {
type Stream = DuplexStream;
fn authority(&self) -> &'static str {
"127.0.0.1:11434"
}
async fn connect(&self) -> std::io::Result<DuplexStream> {
Err(std::io::Error::new(
std::io::ErrorKind::ConnectionRefused,
"nothing is listening",
))
}
}
let (mut client, mut edge) = duplex(64 * 1024);
client.write_all(&authed("GET")).await.unwrap();
client.shutdown().await.unwrap();
let (credential, _) = Credential::new(&supplied()).expect("a usable token");
let outcome = serve_exchange(&mut edge, &credential, &Refusing, TEST_PEER)
.await
.expect("a refused backend is an answer, not a transport failure");
assert_eq!(outcome, Outcome::BadGateway);
drop(edge);
let mut seen = Vec::new();
client.read_to_end(&mut seen).await.unwrap();
let text = String::from_utf8_lossy(&seen);
assert!(text.starts_with("HTTP/1.1 502"), "got: {text}");
}
struct AnswersEarly(Mutex<Option<DuplexStream>>);
impl AnswersEarly {
fn new(response: &'static str) -> Self {
let (mine, mut theirs) = duplex(1024);
tokio::spawn(async move {
let mut buf = [0u8; 256];
let mut seen = Vec::new();
while !seen.windows(4).any(|w| w == b"\r\n\r\n") {
match theirs.read(&mut buf).await {
Ok(0) | Err(_) => return,
Ok(n) => seen.extend_from_slice(&buf[..n]),
}
}
let _ = theirs.write_all(response.as_bytes()).await;
let _ = theirs.flush().await;
std::future::pending::<()>().await;
});
Self(Mutex::new(Some(mine)))
}
fn hanging_up(response: &'static str) -> Self {
let (mine, mut theirs) = duplex(1024);
tokio::spawn(async move {
let mut buf = [0u8; 256];
let mut seen = Vec::new();
while !seen.windows(4).any(|w| w == b"\r\n\r\n") {
match theirs.read(&mut buf).await {
Ok(0) | Err(_) => return,
Ok(n) => seen.extend_from_slice(&buf[..n]),
}
}
let _ = theirs.write_all(response.as_bytes()).await;
let _ = theirs.flush().await;
drop(theirs);
});
Self(Mutex::new(Some(mine)))
}
}
impl Backend for AnswersEarly {
type Stream = DuplexStream;
fn authority(&self) -> &'static str {
"127.0.0.1:11434"
}
async fn connect(&self) -> std::io::Result<DuplexStream> {
Ok(self.0.lock().await.take().expect("connected once"))
}
}
#[tokio::test]
async fn a_backend_that_answers_before_reading_the_body_is_still_heard() {
let backend =
AnswersEarly::new("HTTP/1.1 413 Payload Too Large\r\nContent-Length: 2\r\n\r\nno");
let (mut client, mut edge) = duplex(256 * 1024);
let body = "x".repeat(64 * 1024);
let request = format!(
"POST /v1/chat/completions HTTP/1.1\r\nHost: x\r\n\
Authorization: Bearer {TOKEN}\r\nContent-Length: {}\r\n\r\n{body}",
body.len()
);
client.write_all(request.as_bytes()).await.unwrap();
let (credential, _) = Credential::new(&supplied()).expect("a usable token");
let outcome = tokio::time::timeout(
std::time::Duration::from_secs(5),
serve_exchange(&mut edge, &credential, &backend, TEST_PEER),
)
.await
.expect("the answer is already in hand; waiting on the body is waiting forever")
.expect("no transport failure");
assert_eq!(outcome, Outcome::Forwarded);
drop(edge);
let mut seen = Vec::new();
client.read_to_end(&mut seen).await.unwrap();
let text = String::from_utf8_lossy(&seen);
assert!(
text.starts_with("HTTP/1.1 413"),
"the backend's answer must reach the client: {text}"
);
assert!(text.ends_with("no"), "body included: {text}");
}
#[tokio::test(start_paused = true)]
async fn a_peer_that_never_finishes_asking_is_timed_out() {
let backend = CountingBackend::new(OK_RESPONSE);
let (mut client, mut edge) = duplex(64 * 1024);
client
.write_all(b"GET /v1/models HTTP/1.1\r\nHost: x\r\n")
.await
.unwrap();
let (credential, _) = Credential::new(&supplied()).expect("a usable token");
let outcome = serve_exchange(&mut edge, &credential, &backend, TEST_PEER)
.await
.expect("a timeout is not a transport failure");
assert_eq!(outcome, Outcome::TimedOut);
assert_eq!(backend.connects(), 0, "and the backend never heard of it");
drop(edge);
let mut seen = Vec::new();
client.read_to_end(&mut seen).await.unwrap();
assert!(
seen.is_empty(),
"a peer that never finished asking is owed no answer: {seen:?}"
);
}
#[tokio::test(start_paused = true)]
async fn a_slow_head_that_arrives_in_time_is_served_normally() {
let backend = CountingBackend::new(OK_RESPONSE);
let (mut client, mut edge) = duplex(64 * 1024);
let auth = format!("Bearer {TOKEN}");
tokio::spawn(async move {
client
.write_all(b"GET /v1/models HTTP/1.1\r\nHost: x\r\n")
.await
.unwrap();
tokio::time::sleep(std::time::Duration::from_secs(5)).await;
client
.write_all(format!("Authorization: {auth}\r\n\r\n").as_bytes())
.await
.unwrap();
tokio::time::sleep(std::time::Duration::from_secs(45)).await;
});
let (credential, _) = Credential::new(&supplied()).expect("a usable token");
let outcome = serve_exchange(&mut edge, &credential, &backend, TEST_PEER)
.await
.expect("no transport failure");
assert_eq!(outcome, Outcome::Forwarded);
assert_eq!(backend.connects(), 1);
}
const OK_TEXT: &str = "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\
Content-Length: 2\r\n\r\n{}";
struct ReadsWholeBody {
stream: Mutex<Option<DuplexStream>>,
saw_eof: Arc<AtomicBool>,
}
#[derive(Clone, Copy)]
enum OnEarlyEnd {
HangUp,
Answer(&'static str),
Hold,
}
impl ReadsWholeBody {
fn new(response: &'static str, after_eof: OnEarlyEnd) -> Self {
let (mine, mut theirs) = duplex(256 * 1024);
let saw_eof = Arc::new(AtomicBool::new(false));
let flag = saw_eof.clone();
tokio::spawn(async move {
let mut buf = [0u8; 4096];
let mut seen = Vec::new();
while !seen.windows(4).any(|w| w == b"\r\n\r\n") {
match theirs.read(&mut buf).await {
Ok(0) | Err(_) => return,
Ok(n) => seen.extend_from_slice(&buf[..n]),
}
}
let head_end = seen
.windows(4)
.position(|w| w == b"\r\n\r\n")
.expect("head ends")
+ 4;
while !body_complete(&seen, head_end) {
match theirs.read(&mut buf).await {
Ok(0) => {
flag.store(true, Ordering::SeqCst);
match after_eof {
OnEarlyEnd::HangUp => {}
OnEarlyEnd::Answer(text) => {
let _ = theirs.write_all(text.as_bytes()).await;
let _ = theirs.flush().await;
}
OnEarlyEnd::Hold => std::future::pending::<()>().await,
}
return;
}
Err(_) => return,
Ok(n) => seen.extend_from_slice(&buf[..n]),
}
}
let _ = theirs.write_all(response.as_bytes()).await;
let _ = theirs.flush().await;
});
Self {
stream: Mutex::new(Some(mine)),
saw_eof,
}
}
fn saw_eof(&self) -> bool {
self.saw_eof.load(Ordering::SeqCst)
}
}
impl Backend for ReadsWholeBody {
type Stream = DuplexStream;
fn authority(&self) -> &'static str {
"127.0.0.1:11434"
}
async fn connect(&self) -> std::io::Result<DuplexStream> {
Ok(self.stream.lock().await.take().expect("connected once"))
}
}
fn body_complete(seen: &[u8], head_end: usize) -> bool {
let head = String::from_utf8_lossy(&seen[..head_end]).to_ascii_lowercase();
if head.contains("transfer-encoding: chunked") {
return seen[head_end..].windows(5).any(|w| w == b"0\r\n\r\n");
}
let declared = head
.split("content-length:")
.nth(1)
.and_then(|rest| rest.split("\r\n").next())
.and_then(|value| value.trim().parse::<usize>().ok())
.unwrap_or(0);
seen.len() - head_end >= declared
}
struct Vanishing(Mutex<Option<DuplexStream>>);
impl Vanishing {
fn new() -> Self {
let (mine, mut theirs) = duplex(1024);
tokio::spawn(async move {
let mut buf = [0u8; 256];
let mut seen = Vec::new();
while !seen.windows(4).any(|w| w == b"\r\n\r\n") {
match theirs.read(&mut buf).await {
Ok(0) | Err(_) => return,
Ok(n) => seen.extend_from_slice(&buf[..n]),
}
}
drop(theirs);
});
Self(Mutex::new(Some(mine)))
}
}
impl Backend for Vanishing {
type Stream = DuplexStream;
fn authority(&self) -> &'static str {
"127.0.0.1:11434"
}
async fn connect(&self) -> std::io::Result<DuplexStream> {
Ok(self.0.lock().await.take().expect("connected once"))
}
}
async fn drive<B: Backend + Sync>(
request: &[u8],
backend: &B,
hang_up: bool,
patience: std::time::Duration,
) -> (Outcome, String) {
let (mut client, mut edge) = duplex(256 * 1024);
client.write_all(request).await.unwrap();
if hang_up {
client.shutdown().await.unwrap();
}
let (credential, _) = Credential::new(&supplied()).expect("a usable token");
let outcome = tokio::time::timeout(
patience,
serve_exchange(&mut edge, &credential, backend, TEST_PEER),
)
.await
.expect("the exchange must not hang")
.expect("no transport failure");
drop(edge);
let mut seen = Vec::new();
client.read_to_end(&mut seen).await.unwrap();
(outcome, String::from_utf8_lossy(&seen).into_owned())
}
async fn against<B: Backend + Sync>(
request: &[u8],
backend: &B,
hang_up: bool,
) -> (Outcome, String) {
drive(request, backend, hang_up, std::time::Duration::from_secs(5)).await
}
fn post(declared: usize, body: &str) -> Vec<u8> {
format!(
"POST /v1/chat/completions HTTP/1.1\r\nHost: x\r\n\
Authorization: Bearer {TOKEN}\r\nContent-Length: {declared}\r\n\r\n{body}"
)
.into_bytes()
}
#[tokio::test]
async fn a_truncated_request_body_is_answered_rather_than_waited_on() {
let backend = ReadsWholeBody::new(OK_TEXT, OnEarlyEnd::HangUp);
let (outcome, seen) = against(&post(1000, "{\"model\":\""), &backend, true).await;
assert_eq!(outcome, Outcome::Unfinished);
assert!(seen.starts_with("HTTP/1.1 400"), "got: {seen}");
assert!(seen.contains("incomplete_request"), "got: {seen}");
assert!(
backend.saw_eof(),
"the backend must be told where the body stopped, not merely abandoned"
);
}
#[tokio::test]
async fn a_complete_request_body_reaches_the_same_backend_that_hangs_on_a_short_one() {
let backend = ReadsWholeBody::new(OK_TEXT, OnEarlyEnd::HangUp);
let body = "{\"model\":\"m\"}";
let (outcome, seen) = against(&post(body.len(), body), &backend, true).await;
assert_eq!(outcome, Outcome::Forwarded);
assert!(seen.starts_with("HTTP/1.1 200"), "got: {seen}");
assert!(!backend.saw_eof(), "a complete body is not an early end");
}
#[tokio::test]
async fn a_backend_that_answers_the_short_body_is_heard_rather_than_overridden() {
let backend = ReadsWholeBody::new(
OK_TEXT,
OnEarlyEnd::Answer("HTTP/1.1 400 Bad Request\r\nContent-Length: 9\r\n\r\ntruncated"),
);
let (outcome, seen) = against(&post(1000, "{\"model\":\""), &backend, true).await;
assert_eq!(outcome, Outcome::Forwarded);
assert!(
seen.ends_with("truncated"),
"the backend's own words, not ours: {seen}"
);
}
#[tokio::test]
async fn an_answer_followed_by_a_hangup_is_relayed_rather_than_charged_to_the_client() {
let backend =
AnswersEarly::hanging_up("HTTP/1.1 413 Payload Too Large\r\nContent-Length: 2\r\n\r\nno");
let body = "x".repeat(64 * 1024);
let (outcome, seen) = against(&post(body.len(), &body), &backend, false).await;
assert_eq!(outcome, Outcome::Forwarded);
assert!(seen.starts_with("HTTP/1.1 413"), "got: {seen}");
}
#[tokio::test]
async fn a_backend_that_hangs_up_mid_body_is_a_gateway_failure_rather_than_a_client_one() {
let backend = Vanishing::new();
let body = "x".repeat(64 * 1024);
let (outcome, seen) = against(&post(body.len(), &body), &backend, false).await;
assert_eq!(outcome, Outcome::BadGateway);
assert!(seen.starts_with("HTTP/1.1 502"), "got: {seen}");
}
#[tokio::test]
async fn a_chunked_body_with_an_unreadable_size_is_refused_rather_than_relayed_on() {
let backend = ReadsWholeBody::new(OK_TEXT, OnEarlyEnd::HangUp);
let request = format!(
"POST /v1/chat/completions HTTP/1.1\r\nHost: x\r\n\
Authorization: Bearer {TOKEN}\r\nTransfer-Encoding: chunked\r\n\r\nzz\r\n"
);
let (outcome, seen) = against(request.as_bytes(), &backend, false).await;
assert_eq!(outcome, Outcome::Unfinished);
assert!(seen.starts_with("HTTP/1.1 400"), "got: {seen}");
assert!(backend.saw_eof(), "the backend is told here too");
}
#[tokio::test]
async fn a_well_formed_chunked_body_is_forwarded_rather_than_refused() {
let backend = ReadsWholeBody::new(OK_TEXT, OnEarlyEnd::HangUp);
let request = format!(
"POST /v1/chat/completions HTTP/1.1\r\nHost: x\r\n\
Authorization: Bearer {TOKEN}\r\nTransfer-Encoding: chunked\r\n\r\n\
a\r\n0123456789\r\n0\r\n\r\n"
);
let (outcome, seen) = against(request.as_bytes(), &backend, false).await;
assert_eq!(outcome, Outcome::Forwarded);
assert!(seen.starts_with("HTTP/1.1 200"), "got: {seen}");
}
#[tokio::test(start_paused = true)]
async fn a_backend_that_neither_answers_nor_closes_is_not_waited_on_forever() {
let backend = ReadsWholeBody::new(OK_TEXT, OnEarlyEnd::Hold);
let (outcome, seen) = drive(
&post(1000, "{\"model\":\""),
&backend,
true,
crate::request_body::ANSWER_GRACE * 100,
)
.await;
assert_eq!(outcome, Outcome::Unfinished);
assert!(seen.starts_with("HTTP/1.1 400"), "got: {seen}");
assert!(backend.saw_eof(), "it was told, it simply did nothing");
}
#[tokio::test(start_paused = true)]
async fn a_slow_answer_to_a_complete_body_is_waited_for_however_long_it_takes() {
let backend = Deliberating::new(OK_TEXT, crate::request_body::ANSWER_GRACE * 100);
let body = "{\"model\":\"m\"}";
let (outcome, seen) = drive(
&post(body.len(), body),
&backend,
true,
crate::request_body::ANSWER_GRACE * 1000,
)
.await;
assert_eq!(outcome, Outcome::Forwarded);
assert!(seen.starts_with("HTTP/1.1 200"), "got: {seen}");
}
struct Deliberating(Mutex<Option<DuplexStream>>);
impl Deliberating {
fn new(response: &'static str, think_for: std::time::Duration) -> Self {
let (mine, mut theirs) = duplex(256 * 1024);
tokio::spawn(async move {
let mut buf = [0u8; 4096];
let mut seen = Vec::new();
while !seen.windows(4).any(|w| w == b"\r\n\r\n") {
match theirs.read(&mut buf).await {
Ok(0) | Err(_) => return,
Ok(n) => seen.extend_from_slice(&buf[..n]),
}
}
let head_end = seen
.windows(4)
.position(|w| w == b"\r\n\r\n")
.expect("head ends")
+ 4;
while !body_complete(&seen, head_end) {
match theirs.read(&mut buf).await {
Ok(0) | Err(_) => return,
Ok(n) => seen.extend_from_slice(&buf[..n]),
}
}
tokio::time::sleep(think_for).await;
let _ = theirs.write_all(response.as_bytes()).await;
let _ = theirs.flush().await;
});
Self(Mutex::new(Some(mine)))
}
}
impl Backend for Deliberating {
type Stream = DuplexStream;
fn authority(&self) -> &'static str {
"127.0.0.1:11434"
}
async fn connect(&self) -> std::io::Result<DuplexStream> {
Ok(self.0.lock().await.take().expect("connected once"))
}
}
const LOG_TOKEN: &str = "sk-zzq-tracing-sentinel";
const LOG_QUERY: &str = "sk-zzq-query-sentinel";
thread_local! {
static CAPTURED: RefCell<Option<Vec<u8>>> = const { RefCell::new(None) };
}
#[derive(Clone, Copy)]
struct Sink;
impl std::io::Write for Sink {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
CAPTURED.with(|slot| {
if let Some(into) = slot.borrow_mut().as_mut() {
into.extend_from_slice(buf);
}
});
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
impl tracing_subscriber::fmt::MakeWriter<'_> for Sink {
type Writer = Self;
fn make_writer(&self) -> Self::Writer {
*self
}
}
fn capturing() {
static ONCE: std::sync::Once = std::sync::Once::new();
ONCE.call_once(|| {
let subscriber = tracing_subscriber::fmt()
.with_writer(Sink)
.with_ansi(false)
.with_max_level(tracing::Level::TRACE)
.finish();
tracing::subscriber::set_global_default(subscriber)
.expect("nothing else in this binary may install a subscriber");
});
}
async fn logged(
request: &[u8],
policy: &TokenPolicy,
backend: &CountingBackend,
) -> (Outcome, String) {
capturing();
CAPTURED.with(|slot| *slot.borrow_mut() = Some(Vec::new()));
let (outcome, _) = exchange(request, policy, backend).await;
let written = CAPTURED
.with(|slot| slot.borrow_mut().take())
.expect("armed just above");
(outcome, String::from_utf8(written).expect("utf-8"))
}
fn sensitive_get() -> Vec<u8> {
let mut req =
format!("GET /v1/models?api-key={LOG_QUERY} HTTP/1.1\r\nHost: 127.0.0.1:8080\r\n")
.into_bytes();
req.extend_from_slice(format!("Authorization: Bearer {LOG_TOKEN}\r\n").as_bytes());
req.extend_from_slice(b"\r\n");
req
}
#[tokio::test]
async fn a_forwarded_exchange_reports_what_it_did() {
let backend = CountingBackend::new(OK_RESPONSE);
let (outcome, log) = logged(
&sensitive_get(),
&TokenPolicy::Supplied(LOG_TOKEN.to_owned()),
&backend,
)
.await;
assert_eq!(outcome, Outcome::Forwarded);
for field in [
"method=\"GET\"",
"path=\"/v1/models\"",
"status=200",
"outcome=\"forwarded\"",
"elapsed_ms=",
] {
assert!(log.contains(field), "no {field} in:\n{log}");
}
}
#[tokio::test]
async fn no_line_repeats_the_credential_or_the_query_string() {
let request = sensitive_get();
let raw = String::from_utf8(request.clone()).unwrap();
assert!(raw.contains(LOG_TOKEN), "the request must carry the token");
assert!(raw.contains(LOG_QUERY), "the request must carry the query");
let backend = CountingBackend::new(OK_RESPONSE);
let (outcome, log) = logged(
&request,
&TokenPolicy::Supplied(LOG_TOKEN.to_owned()),
&backend,
)
.await;
assert_eq!(outcome, Outcome::Forwarded);
assert!(!log.contains(LOG_TOKEN), "the token is in the log:\n{log}");
assert!(!log.contains(LOG_QUERY), "the query is in the log:\n{log}");
assert!(
log.contains("path=\"/v1/models\""),
"captured nothing:\n{log}"
);
assert!(
!log.contains("api-key"),
"even the parameter name is gone:\n{log}"
);
}
#[tokio::test]
async fn a_grant_that_admits_is_logged_without_the_grant() {
const GRANT: &str = "sk-zzq-grant-sentinel";
let mut request = b"GET /v1/models HTTP/1.1\r\nHost: 127.0.0.1:8080\r\n".to_vec();
request.extend_from_slice(format!("Authorization: Bearer {GRANT}\r\n\r\n").as_bytes());
let backend = CountingBackend::new(OK_RESPONSE);
let (credential, _) = Credential::new(&supplied()).expect("a usable policy");
assert!(credential.grant(GRANT.to_owned(), std::time::Duration::from_mins(1), None));
capturing();
CAPTURED.with(|slot| *slot.borrow_mut() = Some(Vec::new()));
let (mut client, mut edge) = duplex(64 * 1024);
client.write_all(&request).await.unwrap();
client.shutdown().await.unwrap();
let outcome = serve_exchange(&mut edge, &credential, &backend, TEST_PEER)
.await
.expect("no transport failure");
drop(edge);
let log = CAPTURED
.with(|slot| slot.borrow_mut().take())
.map(|bytes| String::from_utf8(bytes).expect("utf-8"))
.expect("armed just above");
assert_eq!(outcome, Outcome::Forwarded, "the grant admitted");
assert_eq!(backend.connects(), 1);
assert!(
log.contains("outcome=\"forwarded\""),
"captured nothing:\n{log}"
);
assert!(!log.contains(GRANT), "the grant is in the log:\n{log}");
}
#[tokio::test]
async fn a_refused_exchange_names_the_refusal_and_the_request() {
let backend = CountingBackend::new(OK_RESPONSE);
let (outcome, log) = logged(&sensitive_get(), &supplied(), &backend).await;
assert_eq!(outcome, Outcome::Unauthorized);
assert_eq!(backend.connects(), 0, "the backend must not be contacted");
for field in [
"outcome=\"unauthorized\"",
"method=\"GET\"",
"path=\"/v1/models\"",
] {
assert!(log.contains(field), "no {field} in:\n{log}");
}
assert!(!log.contains(LOG_TOKEN), "the token is in the log:\n{log}");
assert!(!log.contains(LOG_QUERY), "the query is in the log:\n{log}");
assert!(
!log.contains("status="),
"there was no upstream status:\n{log}"
);
}
#[tokio::test]
async fn an_exchange_is_unchanged_by_being_watched() {
let watched = CountingBackend::new(OK_RESPONSE);
let (with_capture, log) = logged(
&sensitive_get(),
&TokenPolicy::Supplied(LOG_TOKEN.to_owned()),
&watched,
)
.await;
assert!(
!log.is_empty(),
"the armed run must have captured something"
);
let unwatched = CountingBackend::new(OK_RESPONSE);
let (without, seen) = exchange(
&sensitive_get(),
&TokenPolicy::Supplied(LOG_TOKEN.to_owned()),
&unwatched,
)
.await;
assert_eq!(with_capture, without);
assert_eq!(watched.connects(), unwatched.connects());
assert!(String::from_utf8_lossy(&seen).starts_with("HTTP/1.1 200"));
}
struct Truncating(Mutex<Option<DuplexStream>>);
impl Truncating {
fn new() -> Self {
let (mine, mut theirs) = duplex(1024);
tokio::spawn(async move {
let mut buf = [0u8; 256];
let mut seen = Vec::new();
while !seen.windows(4).any(|w| w == b"\r\n\r\n") {
match theirs.read(&mut buf).await {
Ok(0) | Err(_) => return,
Ok(n) => seen.extend_from_slice(&buf[..n]),
}
}
let _ = theirs
.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 10\r\n\r\nfive!")
.await;
let _ = theirs.flush().await;
drop(theirs);
});
Self(Mutex::new(Some(mine)))
}
}
impl Backend for Truncating {
type Stream = DuplexStream;
fn authority(&self) -> &'static str {
"127.0.0.1:11434"
}
async fn connect(&self) -> std::io::Result<DuplexStream> {
Ok(self.0.lock().await.take().expect("connected once"))
}
}
async fn logged_failure<B: Backend + Sync>(
request: &[u8],
backend: &B,
) -> (std::io::Error, String) {
capturing();
CAPTURED.with(|slot| *slot.borrow_mut() = Some(Vec::new()));
let (mut client, mut edge) = duplex(64 * 1024);
client.write_all(request).await.unwrap();
let (credential, _) = Credential::new(&supplied()).expect("a usable token");
let error = tokio::time::timeout(
std::time::Duration::from_secs(5),
serve_exchange(&mut edge, &credential, backend, TEST_PEER),
)
.await
.expect("the exchange must not hang")
.expect_err("this exchange must fail at the transport");
let written = CAPTURED
.with(|slot| slot.borrow_mut().take())
.expect("armed just above");
(error, String::from_utf8(written).expect("utf-8"))
}
#[tokio::test]
async fn a_backend_that_under_delivers_its_body_is_reported_as_a_failure() {
let backend = Truncating::new();
let (error, log) = logged_failure(&get(Some(&format!("Bearer {TOKEN}"))), &backend).await;
assert_eq!(
error.kind(),
std::io::ErrorKind::UnexpectedEof,
"the body ended before its declared length: {error}"
);
assert!(
log.contains("exchange failed"),
"no failure line in:\n{log}"
);
assert!(log.contains("method=\"GET\""), "unattributable:\n{log}");
assert!(
log.contains("status=200"),
"the backend's head did arrive, and the line should say so:\n{log}"
);
}
#[tokio::test]
async fn a_refused_backend_response_still_reports_the_status_it_sent() {
let backend = CountingBackend::new(
b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nTransfer-Encoding: chunked\r\n\r\n{}",
);
let (outcome, log) = logged(
&get(Some(&format!("Bearer {TOKEN}"))),
&supplied(),
&backend,
)
.await;
assert_eq!(outcome, Outcome::BadGateway);
assert!(log.contains("outcome=\"bad_gateway\""), "{log}");
assert!(
log.contains("status=200"),
"the status it refused is missing:\n{log}"
);
}
fn expecting(value: &str) -> Vec<u8> {
format!(
"GET /v1/models HTTP/1.1\r\nHost: 127.0.0.1:8080\r\n\
Authorization: Bearer {TOKEN}\r\nExpect: {value}\r\n\r\n"
)
.into_bytes()
}
#[tokio::test]
async fn an_expectation_of_continue_is_answered_before_the_response() {
let backend = CountingBackend::new(OK_RESPONSE);
let (outcome, seen) = exchange(&expecting("100-continue"), &supplied(), &backend).await;
assert_eq!(outcome, Outcome::Forwarded);
let text = String::from_utf8(seen).expect("ascii");
assert!(
text.starts_with("HTTP/1.1 100 Continue\r\n\r\n"),
"the interim answer must come first: {text:?}"
);
assert!(
text.contains("HTTP/1.1 200 OK"),
"the final response is still sent: {text:?}"
);
}
#[tokio::test]
async fn a_request_that_expects_nothing_gets_no_interim_response() {
let backend = CountingBackend::new(OK_RESPONSE);
let (outcome, seen) = exchange(
&get(Some(&format!("Bearer {TOKEN}"))),
&supplied(),
&backend,
)
.await;
assert_eq!(outcome, Outcome::Forwarded);
let text = String::from_utf8(seen).expect("ascii");
assert!(text.starts_with("HTTP/1.1 200"), "{text:?}");
assert!(!text.contains("100 Continue"), "{text:?}");
}
#[tokio::test]
async fn an_expectation_is_not_answered_for_a_request_that_is_refused() {
let backend = CountingBackend::new(OK_RESPONSE);
let (outcome, seen) = exchange(&expecting("100-continue"), &supplied(), &backend).await;
assert_eq!(outcome, Outcome::Forwarded, "the sanity case");
let refused = CountingBackend::new(OK_RESPONSE);
let mut req = expecting("100-continue");
req = String::from_utf8(req)
.unwrap()
.replace(&format!("Authorization: Bearer {TOKEN}\r\n"), "")
.into_bytes();
let (outcome, denied) = exchange(&req, &supplied(), &refused).await;
assert_eq!(outcome, Outcome::Unauthorized);
assert_eq!(refused.connects(), 0, "the backend was not contacted");
let text = String::from_utf8(denied).expect("ascii");
assert!(text.starts_with("HTTP/1.1 401"), "{text:?}");
assert!(!text.contains("100 Continue"), "{text:?}");
drop(seen);
}
#[tokio::test]
async fn the_expectation_is_matched_by_name_and_not_by_presence() {
for value in ["100-continue", "100-CONTINUE", "100-Continue"] {
let backend = CountingBackend::new(OK_RESPONSE);
let (_, seen) = exchange(&expecting(value), &supplied(), &backend).await;
let text = String::from_utf8(seen).expect("ascii");
assert!(text.starts_with("HTTP/1.1 100"), "{value}: {text:?}");
}
for value in ["200-ok", "something-else", ""] {
let backend = CountingBackend::new(OK_RESPONSE);
let (_, seen) = exchange(&expecting(value), &supplied(), &backend).await;
let text = String::from_utf8(seen).expect("ascii");
assert!(
!text.contains("100 Continue"),
"{value:?} is not an expectation this edge answers: {text:?}"
);
}
}
#[tokio::test]
async fn a_backend_interim_response_is_not_relayed_on_top_of_the_edge_s() {
let (outcome, seen) = against_keepalive(
&expecting("100-continue"),
"HTTP/1.1 100 Continue\r\n\r\nHTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\n{}",
)
.await;
assert_eq!(outcome, Outcome::Forwarded);
assert_eq!(
seen.matches("100 Continue").count(),
1,
"exactly one interim response reaches the client: {seen:?}"
);
assert!(seen.contains("HTTP/1.1 200 OK"), "{seen:?}");
}
#[tokio::test]
async fn a_backend_that_was_never_reached_says_so_in_the_body() {
struct Refusing;
impl Backend for Refusing {
type Stream = DuplexStream;
fn authority(&self) -> &'static str {
"127.0.0.1:11434"
}
async fn connect(&self) -> std::io::Result<DuplexStream> {
Err(std::io::Error::new(
std::io::ErrorKind::ConnectionRefused,
"nothing is listening",
))
}
}
let (mut client, mut edge) = duplex(64 * 1024);
client.write_all(&authed("GET")).await.unwrap();
client.shutdown().await.unwrap();
let (credential, _) = Credential::new(&supplied()).expect("a usable token");
let outcome = serve_exchange(&mut edge, &credential, &Refusing, TEST_PEER)
.await
.unwrap();
drop(edge);
let mut seen = Vec::new();
client.read_to_end(&mut seen).await.unwrap();
let text = String::from_utf8(seen).expect("ascii");
assert_eq!(outcome, Outcome::BadGateway);
assert!(text.contains(r#""code":"backend_unreachable""#), "{text}");
assert!(!text.contains(r#""code":"bad_gateway""#), "{text}");
}
#[tokio::test]
async fn a_backend_that_answered_unreadably_keeps_the_other_body() {
let backend = CountingBackend::new(
b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nTransfer-Encoding: chunked\r\n\r\n{}",
);
let (outcome, seen) = exchange(
&get(Some(&format!("Bearer {TOKEN}"))),
&supplied(),
&backend,
)
.await;
let text = String::from_utf8(seen).expect("ascii");
assert_eq!(outcome, Outcome::BadGateway);
assert_eq!(backend.connects(), 1, "the backend was reached");
assert!(text.contains(r#""code":"bad_gateway""#), "{text}");
assert!(!text.contains(r#""code":"backend_unreachable""#), "{text}");
}
#[tokio::test]
async fn an_expectation_list_containing_continue_is_still_answered() {
for value in ["100-continue, foo", "foo, 100-continue", "foo,100-CONTINUE"] {
let backend = CountingBackend::new(OK_RESPONSE);
let (_, seen) = exchange(&expecting(value), &supplied(), &backend).await;
let text = String::from_utf8(seen).expect("ascii");
assert!(
text.starts_with("HTTP/1.1 100 Continue"),
"{value:?} asks for the continue: {text:?}"
);
}
for value in ["foo", "foo, bar", "100-continues"] {
let backend = CountingBackend::new(OK_RESPONSE);
let (_, seen) = exchange(&expecting(value), &supplied(), &backend).await;
let text = String::from_utf8(seen).expect("ascii");
assert!(!text.contains("100 Continue"), "{value:?}: {text:?}");
}
}
#[tokio::test]
async fn the_backend_is_told_which_named_token_admitted_and_nothing_when_none_did() {
const LAPTOP: &str = "sk-zzq-laptop-sentinel";
let (credential, _) = Credential::new(&supplied()).expect("a usable policy");
credential
.add_named("laptop", LAPTOP.to_owned())
.expect("a valid name and token");
for (bearer, expected) in [(LAPTOP, Some("laptop")), (TOKEN, None)] {
let backend = CountingBackend::new(OK_RESPONSE);
let mut req = b"GET /v1/models HTTP/1.1\r\nHost: 127.0.0.1:8080\r\n".to_vec();
req.extend_from_slice(format!("Authorization: Bearer {bearer}\r\n").as_bytes());
req.extend_from_slice(b"X-Modelpipe-Device: forged\r\n\r\n");
let (mut client, mut edge) = duplex(64 * 1024);
client.write_all(&req).await.unwrap();
client.shutdown().await.unwrap();
let outcome = serve_exchange(&mut edge, &credential, &backend, TEST_PEER)
.await
.expect("no transport failure");
drop(edge);
assert_eq!(outcome, Outcome::Forwarded);
let sent = String::from_utf8(backend.received().await).expect("ascii");
let lower = sent.to_ascii_lowercase();
assert!(
!lower.contains("forged"),
"the client's claim is gone: {sent}"
);
match expected {
Some(name) => {
assert_eq!(
lower.matches("\r\nx-modelpipe-device:").count(),
1,
"exactly one device marker: {sent}"
);
assert!(
sent.contains(&format!("X-Modelpipe-Device: {name}")),
"and it names the token: {sent}"
);
}
None => assert!(
!lower.contains("x-modelpipe-device"),
"the primary admits with no device marker: {sent}"
),
}
}
}
#[tokio::test]
async fn the_backend_is_handed_the_upstream_bearer_and_never_the_devices() {
const LAPTOP: &str = "sk-zzq-laptop-sentinel";
const UPSTREAM: &str = "sk-zzq-backend-sentinel";
let (credential, _) = Credential::new(&TokenPolicy::Named).expect("a usable policy");
credential
.add_named("laptop", LAPTOP.to_owned())
.expect("a valid name and token");
assert!(credential.set_upstream(Some(UPSTREAM.to_owned())));
let backend = CountingBackend::new(OK_RESPONSE);
let mut req = b"GET /v1/models HTTP/1.1\r\nHost: 127.0.0.1:8080\r\n".to_vec();
req.extend_from_slice(format!("Authorization: Bearer {LAPTOP}\r\n\r\n").as_bytes());
let (mut client, mut edge) = duplex(64 * 1024);
client.write_all(&req).await.unwrap();
client.shutdown().await.unwrap();
let outcome = serve_exchange(&mut edge, &credential, &backend, TEST_PEER)
.await
.expect("no transport failure");
drop(edge);
assert_eq!(outcome, Outcome::Forwarded);
let sent = String::from_utf8(backend.received().await).expect("ascii");
assert!(
sent.contains(&format!("Authorization: Bearer {UPSTREAM}")),
"the edge's bearer: {sent}"
);
assert!(
!sent.contains(LAPTOP),
"the device's key must never reach the backend: {sent}"
);
assert_eq!(
sent.to_ascii_lowercase()
.matches("\r\nauthorization:")
.count(),
1,
"exactly one: {sent}"
);
}