1use std::io::{self, BufRead, BufReader, Read, Write};
25use std::net::TcpStream;
26use std::time::Duration;
27
28pub const MAX_RESPONSE: usize = 8 * 1024 * 1024;
31
32pub trait Stream: Read + Write {}
37impl<T: Read + Write> Stream for T {}
38
39#[derive(Debug, Clone, PartialEq, Eq)]
41pub struct Url {
42 pub scheme: String,
43 pub host: String,
44 pub port: u16,
45 pub path: String,
47}
48
49impl Url {
50 pub fn parse(s: &str) -> Result<Url, String> {
53 let (scheme, rest) = s
54 .split_once("://")
55 .ok_or_else(|| format!("not an absolute URL: {s}"))?;
56 let scheme = scheme.to_ascii_lowercase();
57 let default_port = match scheme.as_str() {
58 "http" => 80,
59 "https" => 443,
60 other => return Err(format!("unsupported scheme: {other}")),
61 };
62 let (authority, path) = match rest.find('/') {
63 Some(i) => (&rest[..i], &rest[i..]),
64 None => (rest, "/"),
65 };
66 if authority.is_empty() {
67 return Err(format!("missing host in URL: {s}"));
68 }
69 let (host, port) = match authority.rsplit_once(':') {
70 Some((h, p)) if p.chars().all(|c| c.is_ascii_digit()) && !p.is_empty() => (
73 h.to_string(),
74 p.parse().map_err(|_| format!("bad port in {s}"))?,
75 ),
76 _ => (authority.to_string(), default_port),
77 };
78 Ok(Url {
79 scheme,
80 host,
81 port,
82 path: path.to_string(),
83 })
84 }
85
86 pub fn is_tls(&self) -> bool {
87 self.scheme == "https"
88 }
89
90 pub fn host_header(&self) -> String {
92 let default = if self.is_tls() { 443 } else { 80 };
93 if self.port == default {
94 self.host.clone()
95 } else {
96 format!("{}:{}", self.host, self.port)
97 }
98 }
99}
100
101#[derive(Debug, Clone)]
103pub struct Response {
104 pub status: u16,
105 pub headers: Vec<(String, String)>,
106 pub body: Vec<u8>,
107}
108
109impl Response {
110 pub fn header(&self, name: &str) -> Option<&str> {
112 let name = name.to_ascii_lowercase();
113 self.headers
114 .iter()
115 .find(|(k, _)| *k == name)
116 .map(|(_, v)| v.as_str())
117 }
118
119 pub fn is_success(&self) -> bool {
120 (200..300).contains(&self.status)
121 }
122
123 pub fn body_str(&self) -> std::borrow::Cow<'_, str> {
124 String::from_utf8_lossy(&self.body)
125 }
126}
127
128pub fn is_loopback_host(host: &str) -> bool {
134 let h = host.trim_start_matches('[').trim_end_matches(']');
135 if h.eq_ignore_ascii_case("localhost") {
136 return true;
137 }
138 h.parse::<std::net::IpAddr>()
139 .map(|ip| ip.is_loopback())
140 .unwrap_or(false)
141}
142
143pub fn connect_tcp(host: &str, port: u16, timeout: Duration) -> io::Result<TcpStream> {
149 use std::net::ToSocketAddrs;
150 let addr = (host, port).to_socket_addrs()?.next().ok_or_else(|| {
151 io::Error::new(
152 io::ErrorKind::NotFound,
153 format!("cannot resolve {host}:{port}"),
154 )
155 })?;
156 let stream = TcpStream::connect_timeout(&addr, timeout)?;
157 stream.set_read_timeout(Some(timeout))?;
158 stream.set_write_timeout(Some(timeout))?;
159 stream.set_nodelay(true).ok();
160 Ok(stream)
161}
162
163pub fn send<S: Read + Write + ?Sized>(
167 stream: &mut S,
168 host_header: &str,
169 method: &str,
170 path: &str,
171 headers: &[(&str, &str)],
172 body: &[u8],
173) -> io::Result<Response> {
174 let mut req: Vec<u8> = Vec::with_capacity(256 + body.len());
181 if path.contains(['\r', '\n']) {
182 return Err(io::Error::new(
183 io::ErrorKind::InvalidInput,
184 "CR/LF in request target",
185 ));
186 }
187 write!(req, "{method} {path} HTTP/1.1\r\n")?;
188 write!(req, "Host: {host_header}\r\n")?;
189 req.extend_from_slice(b"Connection: close\r\n");
190 for (k, v) in headers {
191 if k.contains(['\r', '\n']) || v.contains(['\r', '\n']) {
194 return Err(io::Error::new(
195 io::ErrorKind::InvalidInput,
196 "CR/LF in header",
197 ));
198 }
199 write!(req, "{k}: {v}\r\n")?;
200 }
201 write!(req, "Content-Length: {}\r\n\r\n", body.len())?;
202 req.extend_from_slice(body);
203 stream.write_all(&req)?;
204 stream.flush()?;
205
206 let mut reader = BufReader::new(stream);
207 read_response(&mut reader)
208}
209
210fn read_head<R: BufRead>(r: &mut R) -> io::Result<(u16, Vec<(String, String)>)> {
213 let mut status_line = String::new();
214 r.read_line(&mut status_line)?;
215 let status = parse_status(&status_line)?;
216
217 let mut headers = Vec::new();
218 loop {
219 let mut line = String::new();
220 if r.read_line(&mut line)? == 0 {
221 break;
222 }
223 let line = line.trim_end_matches(['\r', '\n']);
224 if line.is_empty() {
225 break;
226 }
227 if let Some((k, v)) = line.split_once(':') {
228 headers.push((k.trim().to_ascii_lowercase(), v.trim().to_string()));
229 }
230 }
231 Ok((status, headers))
232}
233
234fn read_response<R: BufRead>(r: &mut R) -> io::Result<Response> {
235 let (status, headers) = read_head(r)?;
236
237 let content_length = headers
238 .iter()
239 .find(|(k, _)| k == "content-length")
240 .and_then(|(_, v)| v.parse::<usize>().ok());
241 let chunked = headers
242 .iter()
243 .any(|(k, v)| k == "transfer-encoding" && v.to_ascii_lowercase().contains("chunked"));
244
245 let body = if chunked {
246 read_chunked(r)?
247 } else if let Some(n) = content_length {
248 read_exact_capped(r, n)?
249 } else {
250 read_to_end_capped(r)?
252 };
253
254 Ok(Response {
255 status,
256 headers,
257 body,
258 })
259}
260
261fn parse_status(line: &str) -> io::Result<u16> {
262 line.split_whitespace()
264 .nth(1)
265 .and_then(|s| s.parse().ok())
266 .ok_or_else(|| {
267 io::Error::new(
268 io::ErrorKind::InvalidData,
269 format!("bad status line: {line:?}"),
270 )
271 })
272}
273
274fn read_exact_capped<R: Read>(r: &mut R, n: usize) -> io::Result<Vec<u8>> {
275 if n > MAX_RESPONSE {
276 return Err(io::Error::new(
277 io::ErrorKind::InvalidData,
278 "response exceeds cap",
279 ));
280 }
281 let mut buf = vec![0u8; n];
282 r.read_exact(&mut buf)?;
283 Ok(buf)
284}
285
286fn read_to_end_capped<R: Read>(r: &mut R) -> io::Result<Vec<u8>> {
287 let mut buf = Vec::new();
288 r.take(MAX_RESPONSE as u64 + 1).read_to_end(&mut buf)?;
289 if buf.len() > MAX_RESPONSE {
290 return Err(io::Error::new(
291 io::ErrorKind::InvalidData,
292 "response exceeds cap",
293 ));
294 }
295 Ok(buf)
296}
297
298fn read_chunked<R: BufRead>(r: &mut R) -> io::Result<Vec<u8>> {
299 let mut body = Vec::new();
300 loop {
301 let mut size_line = String::new();
302 r.read_line(&mut size_line)?;
303 let size_hex = size_line.trim_end_matches(['\r', '\n']);
304 let size_hex = size_hex.split(';').next().unwrap_or("").trim();
306 let size = usize::from_str_radix(size_hex, 16)
307 .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "bad chunk size"))?;
308 if size == 0 {
309 loop {
311 let mut t = String::new();
312 if r.read_line(&mut t)? == 0 || t.trim_end_matches(['\r', '\n']).is_empty() {
313 break;
314 }
315 }
316 break;
317 }
318 if size > MAX_RESPONSE || body.len().saturating_add(size) > MAX_RESPONSE {
329 return Err(io::Error::new(
330 io::ErrorKind::InvalidData,
331 "chunked response exceeds cap",
332 ));
333 }
334 let before = body.len();
340 r.by_ref().take(size as u64).read_to_end(&mut body)?;
341 if body.len() - before != size {
342 return Err(io::Error::new(
343 io::ErrorKind::UnexpectedEof,
344 "truncated chunk",
345 ));
346 }
347 let mut crlf = [0u8; 2];
349 r.read_exact(&mut crlf)?;
350 }
351 Ok(body)
352}
353
354pub struct StreamingResponse<S: Read + Write> {
359 pub status: u16,
360 pub headers: Vec<(String, String)>,
361 reader: BufReader<S>,
362}
363
364impl<S: Read + Write> StreamingResponse<S> {
365 pub fn header(&self, name: &str) -> Option<&str> {
366 let name = name.to_ascii_lowercase();
367 self.headers
368 .iter()
369 .find(|(k, _)| *k == name)
370 .map(|(_, v)| v.as_str())
371 }
372 pub fn is_success(&self) -> bool {
373 (200..300).contains(&self.status)
374 }
375 pub fn content_type(&self) -> Option<&str> {
377 self.header("content-type")
378 .map(|v| v.split(';').next().unwrap_or(v).trim())
379 }
380 pub fn is_event_stream(&self) -> bool {
382 self.content_type() == Some("text/event-stream")
383 }
384 pub fn into_body(mut self) -> io::Result<Vec<u8>> {
386 let content_length = self
387 .headers
388 .iter()
389 .find(|(k, _)| k == "content-length")
390 .and_then(|(_, v)| v.parse::<usize>().ok());
391 let chunked = self
392 .headers
393 .iter()
394 .any(|(k, v)| k == "transfer-encoding" && v.to_ascii_lowercase().contains("chunked"));
395 if chunked {
396 read_chunked(&mut self.reader)
397 } else if let Some(n) = content_length {
398 read_exact_capped(&mut self.reader, n)
399 } else {
400 read_to_end_capped(&mut self.reader)
401 }
402 }
403 pub fn sse(self) -> SseReader<BufReader<S>> {
405 SseReader::new(self.reader)
406 }
407
408 pub fn into_reader(self) -> BufReader<S> {
411 self.reader
412 }
413}
414
415pub fn send_streaming<S: Read + Write>(
420 mut stream: S,
421 host_header: &str,
422 method: &str,
423 path: &str,
424 headers: &[(&str, &str)],
425 body: &[u8],
426) -> io::Result<StreamingResponse<S>> {
427 let mut req: Vec<u8> = Vec::with_capacity(256 + body.len());
430 if path.contains(['\r', '\n']) {
431 return Err(io::Error::new(
432 io::ErrorKind::InvalidInput,
433 "CR/LF in request target",
434 ));
435 }
436 write!(req, "{method} {path} HTTP/1.1\r\n")?;
437 write!(req, "Host: {host_header}\r\n")?;
438 req.extend_from_slice(b"Connection: close\r\n");
439 for (k, v) in headers {
440 if k.contains(['\r', '\n']) || v.contains(['\r', '\n']) {
441 return Err(io::Error::new(
442 io::ErrorKind::InvalidInput,
443 "CR/LF in header",
444 ));
445 }
446 write!(req, "{k}: {v}\r\n")?;
447 }
448 write!(req, "Content-Length: {}\r\n\r\n", body.len())?;
449 req.extend_from_slice(body);
450 stream.write_all(&req)?;
451 stream.flush()?;
452
453 let mut reader = BufReader::new(stream);
454 let (status, headers) = read_head(&mut reader)?;
455 Ok(StreamingResponse {
456 status,
457 headers,
458 reader,
459 })
460}
461
462#[derive(Debug, Clone, Default, PartialEq, Eq)]
465pub struct SseEvent {
466 pub event: Option<String>,
467 pub data: String,
468 pub id: Option<String>,
469}
470
471pub struct SseReader<R: BufRead> {
476 r: R,
477}
478
479impl<R: BufRead> SseReader<R> {
480 pub fn new(r: R) -> SseReader<R> {
481 SseReader { r }
482 }
483
484 pub fn next_event(&mut self) -> io::Result<Option<SseEvent>> {
487 let mut ev = SseEvent::default();
488 let mut saw_field = false;
489 let mut total = 0usize;
490 loop {
491 let mut line = String::new();
492 let n = self.r.read_line(&mut line)?;
493 if n == 0 {
494 return Ok(if saw_field { Some(ev) } else { None });
496 }
497 total += n;
498 if total > MAX_RESPONSE {
499 return Err(io::Error::new(
500 io::ErrorKind::InvalidData,
501 "SSE event exceeds cap",
502 ));
503 }
504 let line = line.trim_end_matches(['\r', '\n']);
505 if line.is_empty() {
506 if saw_field {
508 return Ok(Some(ev));
509 }
510 continue; }
512 if line.starts_with(':') {
513 continue; }
515 let (field, value) = match line.split_once(':') {
516 Some((f, v)) => (f, v.strip_prefix(' ').unwrap_or(v)),
517 None => (line, ""), };
519 saw_field = true;
520 match field {
521 "event" => ev.event = Some(value.to_string()),
522 "id" => ev.id = Some(value.to_string()),
523 "data" => {
524 if !ev.data.is_empty() {
525 ev.data.push('\n');
526 }
527 ev.data.push_str(value);
528 }
529 _ => {} }
531 }
532 }
533}
534
535#[cfg(test)]
536mod tests {
537 use super::*;
538 use std::io::Cursor;
539
540 #[test]
541 fn url_parse_https_default_port() {
542 let u = Url::parse("https://api.openai.com/v1/chat/completions").unwrap();
543 assert_eq!(u.scheme, "https");
544 assert_eq!(u.host, "api.openai.com");
545 assert_eq!(u.port, 443);
546 assert_eq!(u.path, "/v1/chat/completions");
547 assert_eq!(u.host_header(), "api.openai.com");
548 assert!(u.is_tls());
549 }
550
551 #[test]
552 fn url_parse_http_with_port_and_no_path() {
553 let u = Url::parse("http://localhost:8080").unwrap();
554 assert_eq!(u.port, 8080);
555 assert_eq!(u.path, "/");
556 assert_eq!(u.host_header(), "localhost:8080");
557 assert!(!u.is_tls());
558 }
559
560 #[test]
561 fn url_rejects_bad_scheme() {
562 assert!(Url::parse("ftp://x/").is_err());
563 assert!(Url::parse("no-scheme").is_err());
564 }
565
566 #[test]
567 fn response_content_length() {
568 let raw = "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 13\r\n\r\n{\"ok\":true}!!";
569 let mut cur = Cursor::new(raw.as_bytes().to_vec());
570 let resp = read_response(&mut cur).unwrap();
571 assert_eq!(resp.status, 200);
572 assert_eq!(resp.header("content-type"), Some("application/json"));
573 assert_eq!(resp.body, b"{\"ok\":true}!!");
574 assert!(resp.is_success());
575 }
576
577 #[test]
578 fn response_chunked() {
579 let raw = "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n5\r\nhello\r\n6\r\n world\r\n0\r\n\r\n";
580 let mut cur = Cursor::new(raw.as_bytes().to_vec());
581 let resp = read_response(&mut cur).unwrap();
582 assert_eq!(resp.body, b"hello world");
583 }
584
585 #[test]
586 fn cr_lf_header_injection_rejected() {
587 let mut sink: Vec<u8> = Vec::new();
588 let mut stream = Cursor::new(Vec::new());
590 let _ = &mut sink;
591 let err = send(&mut stream, "h", "POST", "/", &[("X", "a\r\nEvil: 1")], b"").unwrap_err();
592 assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
593 }
594
595 struct FakeStream {
598 resp: Cursor<Vec<u8>>,
599 sink: Vec<u8>,
600 }
601 impl Read for FakeStream {
602 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
603 self.resp.read(buf)
604 }
605 }
606 impl Write for FakeStream {
607 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
608 self.sink.extend_from_slice(buf);
609 Ok(buf.len())
610 }
611 fn flush(&mut self) -> io::Result<()> {
612 Ok(())
613 }
614 }
615
616 fn sse_events(body: &str) -> Vec<SseEvent> {
617 let mut r = SseReader::new(BufReader::new(Cursor::new(body.as_bytes().to_vec())));
618 let mut out = Vec::new();
619 while let Some(e) = r.next_event().unwrap() {
620 out.push(e);
621 }
622 out
623 }
624
625 #[test]
626 fn sse_parses_events_with_event_id_and_data() {
627 let body = "event: message\nid: 7\ndata: {\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{}}\n\n";
628 let evs = sse_events(body);
629 assert_eq!(evs.len(), 1);
630 assert_eq!(evs[0].event.as_deref(), Some("message"));
631 assert_eq!(evs[0].id.as_deref(), Some("7"));
632 assert_eq!(evs[0].data, "{\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{}}");
633 }
634
635 #[test]
636 fn sse_joins_multi_data_lines_and_ignores_comments() {
637 let body = ": keep-alive\ndata: line1\ndata: line2\n\ndata: second\n\n";
639 let evs = sse_events(body);
640 assert_eq!(evs.len(), 2);
641 assert_eq!(evs[0].data, "line1\nline2");
642 assert_eq!(evs[1].data, "second");
643 }
644
645 #[test]
646 fn sse_flushes_trailing_event_without_final_blank_line() {
647 let evs = sse_events("data: only\n");
648 assert_eq!(evs.len(), 1);
649 assert_eq!(evs[0].data, "only");
650 }
651
652 #[test]
653 fn send_streaming_reads_head_then_buffers_json_body() {
654 let raw = "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nMcp-Session-Id: abc123\r\nContent-Length: 11\r\n\r\n{\"ok\":true}".to_string();
655 let stream = FakeStream {
656 resp: Cursor::new(raw.into_bytes()),
657 sink: Vec::new(),
658 };
659 let resp = send_streaming(
660 stream,
661 "h",
662 "POST",
663 "/mcp",
664 &[("Accept", "application/json")],
665 b"{}",
666 )
667 .unwrap();
668 assert_eq!(resp.status, 200);
669 assert!(resp.is_success());
670 assert_eq!(resp.content_type(), Some("application/json"));
671 assert!(!resp.is_event_stream());
672 assert_eq!(resp.header("mcp-session-id"), Some("abc123"));
673 assert_eq!(resp.into_body().unwrap(), b"{\"ok\":true}");
674 }
675
676 #[test]
677 fn send_streaming_detects_event_stream_and_pumps_sse() {
678 let raw = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\r\ndata: {\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{\"x\":1}}\n\n".to_string();
679 let stream = FakeStream {
680 resp: Cursor::new(raw.into_bytes()),
681 sink: Vec::new(),
682 };
683 let resp = send_streaming(stream, "h", "POST", "/mcp", &[], b"{}").unwrap();
684 assert!(resp.is_event_stream());
685 let mut sse = resp.sse();
686 let ev = sse.next_event().unwrap().expect("one event");
687 assert!(ev.data.contains("\"id\":1"));
688 assert!(sse.next_event().unwrap().is_none());
689 }
690}