use std::collections::HashMap;
use std::io;
use std::sync::Mutex;
use std::time::Duration;
use bytes::Bytes;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
const MAX_HEAD: usize = 64 * 1024;
const MAX_BODY: usize = 16 << 20;
#[derive(Debug)]
#[non_exhaustive]
pub enum Error {
Url {
url: String,
reason: String,
},
Timeout(&'static str),
Io(io::Error),
Protocol(String),
TooLarge,
}
impl std::fmt::Display for Error {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Url { url, reason } => write!(f, "{url}: {reason}"),
Self::Timeout(what) => write!(f, "{what} timed out"),
Self::Io(e) => std::fmt::Display::fmt(e, f),
Self::Protocol(m) => f.write_str(m),
Self::TooLarge => f.write_str("response body too large"),
}
}
}
impl std::error::Error for Error {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Io(e) => e.source(),
_ => None,
}
}
}
impl From<io::Error> for Error {
fn from(e: io::Error) -> Self {
Self::Io(e)
}
}
pub type Result<T, E = Error> = std::result::Result<T, E>;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Url {
pub host: String,
pub port: u16,
pub path: String,
}
impl Url {
pub fn parse(url: &str) -> Result<Self> {
let invalid = |reason: String| Error::Url {
url: url.to_string(),
reason,
};
let rest = url.strip_prefix("http://").ok_or_else(|| {
invalid("only http:// URLs are supported by the nerve client".to_string())
})?;
let (authority, path) = match rest.find('/') {
Some(i) => (&rest[..i], &rest[i..]),
None => (rest, "/"),
};
let (host, port) = match authority.rsplit_once(':') {
Some((host, port)) if !host.contains(']') || host.ends_with(']') => (
host,
port.parse::<u16>()
.map_err(|_| invalid(format!("bad port '{port}'")))?,
),
_ => (authority, 80),
};
if host.is_empty() {
return Err(invalid("no host".to_string()));
}
Ok(Self {
host: host.trim_matches(['[', ']']).to_string(),
port,
path: path.to_string(),
})
}
fn key(&self) -> String {
format!("{}:{}", self.host, self.port)
}
fn host_header(&self) -> String {
if self.port == 80 {
self.host.clone()
} else {
format!("{}:{}", self.host, self.port)
}
}
}
#[derive(Debug)]
pub struct Response {
pub status: u16,
pub headers: Vec<(String, String)>,
pub body: Bytes,
}
impl Response {
pub fn header(&self, name: &str) -> Option<&str> {
self.headers
.iter()
.find(|(n, _)| n.eq_ignore_ascii_case(name))
.map(|(_, v)| v.as_str())
}
}
enum Framing {
Length(usize),
Chunked,
Close,
}
pub struct Streaming {
pub status: u16,
pub headers: Vec<(String, String)>,
stream: TcpStream,
framing: Framing,
buffered: Vec<u8>,
chunk_left: Option<usize>,
done: bool,
}
impl Streaming {
pub async fn next(&mut self) -> Result<Option<Bytes>> {
loop {
if self.done {
return Ok(None);
}
match self.framing {
Framing::Length(left) => {
if left == 0 {
self.done = true;
return Ok(None);
}
if !self.buffered.is_empty() {
let n = self.buffered.len().min(left);
let out: Vec<u8> = self.buffered.drain(..n).collect();
self.framing = Framing::Length(left - n);
return Ok(Some(Bytes::from(out)));
}
let mut buf = vec![0u8; left.min(16 * 1024)];
let n = self.stream.read(&mut buf).await?;
if n == 0 {
return Err(io::Error::from(io::ErrorKind::UnexpectedEof).into());
}
buf.truncate(n);
self.framing = Framing::Length(left - n);
return Ok(Some(Bytes::from(buf)));
}
Framing::Close => {
if !self.buffered.is_empty() {
return Ok(Some(Bytes::from(std::mem::take(&mut self.buffered))));
}
let mut buf = vec![0u8; 16 * 1024];
let n = self.stream.read(&mut buf).await?;
if n == 0 {
self.done = true;
return Ok(None);
}
buf.truncate(n);
return Ok(Some(Bytes::from(buf)));
}
Framing::Chunked => {
if let Some(piece) = self.next_chunk_piece().await? {
if piece.is_empty() {
continue;
}
return Ok(Some(piece));
}
self.done = true;
return Ok(None);
}
}
}
}
async fn next_chunk_piece(&mut self) -> Result<Option<Bytes>> {
match self.chunk_left {
None => {
let line = self.read_line().await?;
let size = line.split(';').next().unwrap_or("").trim();
let size = usize::from_str_radix(size, 16)
.map_err(|_| Error::Protocol(format!("bad chunk size '{size}'")))?;
if size == 0 {
while !self.read_line().await?.is_empty() {}
return Ok(None);
}
self.chunk_left = Some(size);
Ok(Some(Bytes::new()))
}
Some(0) => {
self.read_line().await?;
self.chunk_left = None;
Ok(Some(Bytes::new()))
}
Some(left) => {
if self.buffered.is_empty() {
let mut buf = vec![0u8; left.min(16 * 1024)];
let n = self.stream.read(&mut buf).await?;
if n == 0 {
return Err(io::Error::from(io::ErrorKind::UnexpectedEof).into());
}
buf.truncate(n);
self.buffered = buf;
}
let n = self.buffered.len().min(left);
let out: Vec<u8> = self.buffered.drain(..n).collect();
self.chunk_left = Some(left - n);
Ok(Some(Bytes::from(out)))
}
}
}
async fn read_line(&mut self) -> Result<String> {
loop {
if let Some(i) = self.buffered.iter().position(|&b| b == b'\n') {
let line: Vec<u8> = self.buffered.drain(..=i).collect();
let line = String::from_utf8_lossy(&line).trim_end().to_string();
return Ok(line);
}
if self.buffered.len() > MAX_HEAD {
return Err(Error::Protocol("line too long".to_string()));
}
let mut buf = [0u8; 4096];
let n = self.stream.read(&mut buf).await?;
if n == 0 {
return Err(io::Error::from(io::ErrorKind::UnexpectedEof).into());
}
self.buffered.extend_from_slice(&buf[..n]);
}
}
fn keep_alive(&self) -> bool {
!matches!(self.framing, Framing::Close)
&& !self.headers.iter().any(|(n, v)| {
n.eq_ignore_ascii_case("connection") && v.eq_ignore_ascii_case("close")
})
}
}
pub struct Client {
pool: Mutex<HashMap<String, Vec<TcpStream>>>,
timeout: Duration,
}
impl Default for Client {
fn default() -> Self {
Self::new(Duration::from_secs(5))
}
}
impl Client {
pub fn new(timeout: Duration) -> Self {
Self {
pool: Mutex::new(HashMap::new()),
timeout,
}
}
pub async fn warm(&self, url: &Url) -> Result<()> {
let stream = self.connect(url).await?;
self.put_back(url, stream);
Ok(())
}
async fn connect(&self, url: &Url) -> Result<TcpStream> {
let stream = tokio::time::timeout(
self.timeout,
TcpStream::connect((url.host.as_str(), url.port)),
)
.await
.map_err(|_| Error::Timeout("connect"))??;
stream.set_nodelay(true)?;
Ok(stream)
}
fn take(&self, url: &Url) -> Option<TcpStream> {
self.pool
.lock()
.unwrap_or_else(|e| e.into_inner())
.get_mut(&url.key())?
.pop()
}
fn put_back(&self, url: &Url, stream: TcpStream) {
let mut pool = self.pool.lock().unwrap_or_else(|e| e.into_inner());
let idle = pool.entry(url.key()).or_default();
if idle.len() < 8 {
idle.push(stream);
}
}
pub async fn request(
&self,
method: &str,
url: &Url,
headers: &[(&str, &str)],
body: &[u8],
) -> Result<Response> {
tokio::time::timeout(self.timeout, self.request_inner(method, url, headers, body))
.await
.map_err(|_| Error::Timeout("request"))?
}
async fn request_inner(
&self,
method: &str,
url: &Url,
headers: &[(&str, &str)],
body: &[u8],
) -> Result<Response> {
let mut streaming = self.send_inner(method, url, headers, body).await?;
let mut collected = Vec::new();
while let Some(piece) = streaming.next().await? {
collected.extend_from_slice(&piece);
if collected.len() > MAX_BODY {
return Err(Error::TooLarge);
}
}
if streaming.keep_alive() {
self.put_back(url, streaming.stream);
}
Ok(Response {
status: streaming.status,
headers: streaming.headers,
body: Bytes::from(collected),
})
}
pub async fn send(
&self,
method: &str,
url: &Url,
headers: &[(&str, &str)],
body: &[u8],
) -> Result<Streaming> {
tokio::time::timeout(self.timeout, self.send_inner(method, url, headers, body))
.await
.map_err(|_| Error::Timeout("request"))?
}
async fn send_inner(
&self,
method: &str,
url: &Url,
headers: &[(&str, &str)],
body: &[u8],
) -> Result<Streaming> {
let mut head = format!(
"{method} {} HTTP/1.1\r\nHost: {}\r\nUser-Agent: furnace-nerve\r\n",
url.path,
url.host_header()
);
for (name, value) in headers {
head.push_str(name);
head.push_str(": ");
head.push_str(value);
head.push_str("\r\n");
}
if !body.is_empty() || method == "POST" || method == "PUT" {
head.push_str(&format!("Content-Length: {}\r\n", body.len()));
}
head.push_str("\r\n");
let mut request = head.into_bytes();
request.extend_from_slice(body);
if let Some(stream) = self.take(url)
&& let Ok(streaming) = write_and_read_head(stream, &request).await
{
return Ok(streaming);
}
let stream = self.connect(url).await?;
write_and_read_head(stream, &request).await
}
}
async fn write_and_read_head(mut stream: TcpStream, request: &[u8]) -> Result<Streaming> {
stream.write_all(request).await?;
let mut buffered = Vec::with_capacity(4096);
let head_end = loop {
if let Some(i) = find_head_end(&buffered) {
break i;
}
if buffered.len() > MAX_HEAD {
return Err(Error::Protocol("response head too large".to_string()));
}
let mut buf = [0u8; 4096];
let n = stream.read(&mut buf).await?;
if n == 0 {
return Err(io::Error::from(io::ErrorKind::UnexpectedEof).into());
}
buffered.extend_from_slice(&buf[..n]);
};
let head = String::from_utf8_lossy(&buffered[..head_end]).to_string();
let rest = buffered.split_off(head_end + 4);
let mut lines = head.split("\r\n");
let status_line = lines.next().unwrap_or("");
let status = status_line
.split_whitespace()
.nth(1)
.and_then(|s| s.parse::<u16>().ok())
.ok_or_else(|| Error::Protocol(format!("bad status line '{status_line}'")))?;
let headers: Vec<(String, String)> = lines
.filter_map(|line| {
let (name, value) = line.split_once(':')?;
Some((name.trim().to_string(), value.trim().to_string()))
})
.collect();
let header = |name: &str| {
headers
.iter()
.find(|(n, _)| n.eq_ignore_ascii_case(name))
.map(|(_, v)| v.as_str())
};
let framing = if header("transfer-encoding")
.is_some_and(|v| v.to_ascii_lowercase().contains("chunked"))
{
Framing::Chunked
} else if let Some(length) = header("content-length").and_then(|v| v.parse::<usize>().ok()) {
Framing::Length(length)
} else if status == 204 || status == 304 || (100..200).contains(&status) {
Framing::Length(0)
} else {
Framing::Close
};
Ok(Streaming {
status,
headers,
stream,
framing,
buffered: rest,
chunk_left: None,
done: false,
})
}
fn find_head_end(buffer: &[u8]) -> Option<usize> {
buffer.windows(4).position(|w| w == b"\r\n\r\n")
}
#[cfg(test)]
mod tests {
use super::*;
use axum::Router;
use axum::body::Body;
use axum::routing::{get, post};
#[test]
fn urls_parse() {
assert_eq!(
Url::parse("http://localhost:9000/hook?x=1").unwrap(),
Url {
host: "localhost".into(),
port: 9000,
path: "/hook?x=1".into()
}
);
assert_eq!(Url::parse("http://example.com").unwrap().path, "/");
assert_eq!(Url::parse("http://example.com").unwrap().port, 80);
assert!(Url::parse("https://example.com").is_err());
assert!(Url::parse("http://:80/").is_err());
}
async fn serve(app: Router) -> Url {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
Url::parse(&format!("http://127.0.0.1:{port}/")).unwrap()
}
#[tokio::test]
async fn posts_and_reuses_the_connection() {
let base =
serve(Router::new().route("/echo", post(|body: Bytes| async move { body }))).await;
let url = Url {
path: "/echo".into(),
..base
};
let client = Client::default();
for i in 0..3 {
let body = format!("hello {i}");
let response = client
.request(
"POST",
&url,
&[("content-type", "text/plain")],
body.as_bytes(),
)
.await
.unwrap();
assert_eq!(response.status, 200);
assert_eq!(response.body, Bytes::from(body));
}
assert_eq!(client.pool.lock().unwrap()[&url.key()].len(), 1);
}
#[tokio::test]
async fn reads_chunked_streams() {
let base = serve(Router::new().route(
"/stream",
get(|| async {
let stream = futures_util_free_stream();
axum::response::Response::builder()
.header("content-type", "text/event-stream")
.body(Body::from_stream(stream))
.unwrap()
}),
))
.await;
let url = Url {
path: "/stream".into(),
..base
};
let client = Client::default();
let mut streaming = client.send("GET", &url, &[], b"").await.unwrap();
assert_eq!(streaming.status, 200);
let mut all = Vec::new();
while let Some(piece) = streaming.next().await.unwrap() {
all.extend_from_slice(&piece);
}
assert_eq!(
String::from_utf8(all).unwrap(),
"data: 0\n\ndata: 1\n\ndata: 2\n\n"
);
}
async fn raw(response: &'static [u8], close: bool) -> Url {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
tokio::spawn(async move {
loop {
let (mut stream, _) = listener.accept().await.unwrap();
tokio::spawn(async move {
let mut request = [0u8; 4096];
let _ = stream.read(&mut request).await;
let _ = stream.write_all(response).await;
if !close {
tokio::time::sleep(Duration::from_secs(60)).await;
}
});
}
});
Url::parse(&format!("http://127.0.0.1:{port}/")).unwrap()
}
#[tokio::test]
async fn a_body_that_stalls_times_out() {
let url = raw(b"HTTP/1.1 200 OK\r\nContent-Length: 100\r\n\r\nhalf", false).await;
let client = Client::new(Duration::from_millis(100));
let started = std::time::Instant::now();
let err = client.request("POST", &url, &[], b"x").await.unwrap_err();
assert!(matches!(err, Error::Timeout(_)), "{err}");
assert!(started.elapsed() < Duration::from_secs(5));
}
#[tokio::test]
async fn a_body_without_length_runs_until_the_connection_closes() {
let url = raw(b"HTTP/1.1 200 OK\r\n\r\nuntil the end", true).await;
let response = Client::default()
.request("GET", &url, &[], b"")
.await
.unwrap();
assert_eq!(response.status, 200);
assert_eq!(&response.body[..], b"until the end");
}
#[tokio::test]
async fn a_truncated_head_is_an_error() {
let url = raw(b"HTTP/1.1 200 OK\r\nContent-Le", true).await;
let err = Client::default()
.request("GET", &url, &[], b"")
.await
.unwrap_err();
assert!(
matches!(&err, Error::Io(e) if e.kind() == io::ErrorKind::UnexpectedEof),
"{err}"
);
}
#[tokio::test]
async fn a_body_over_the_limit_is_refused() {
let base =
serve(Router::new().route("/big", get(|| async { vec![0u8; MAX_BODY + 1] }))).await;
let url = Url {
path: "/big".into(),
..base
};
let err = Client::default()
.request("GET", &url, &[], b"")
.await
.unwrap_err();
assert!(matches!(err, Error::TooLarge), "{err}");
}
fn futures_util_free_stream()
-> impl tokio_stream::Stream<Item = Result<Bytes, std::convert::Infallible>> {
tokio_stream::iter((0..3).map(|i| Ok(Bytes::from(format!("data: {i}\n\n")))))
}
}