1use sim_cancel::{Cancellation, CancellationReason};
4use sim_kernel::{Error, Result};
5use std::{io, sync::Mutex, time::Duration};
6
7pub type Header = (String, String);
11
12#[derive(Clone, Debug, Eq, PartialEq)]
14pub struct RequestHead {
15 pub method: String,
17 pub target: String,
19 pub headers: Vec<Header>,
21 pub peer: Option<String>,
23 pub local: Option<String>,
25}
26
27#[derive(Clone, Debug, Eq, PartialEq)]
29pub struct ResponseHead {
30 pub status: u16,
32 pub headers: Vec<Header>,
34}
35
36#[derive(Clone, Copy, Debug, Eq, PartialEq)]
38pub struct BodyLimits {
39 pub max_request_bytes: usize,
41 pub max_chunk_bytes: usize,
43}
44
45impl BodyLimits {
46 fn validate(self) -> Result<Self> {
47 if self.max_request_bytes == 0 || self.max_chunk_bytes == 0 {
48 return Err(Error::Eval("raw HTTP body limits must be non-zero".into()));
49 }
50 Ok(self)
51 }
52}
53
54#[derive(Clone, Copy, Debug, Eq, PartialEq)]
56pub enum TrailersPolicy {
57 Deny,
59 Allow,
61}
62
63#[derive(Clone, Debug)]
65pub struct RequestScope {
66 cancellation: Cancellation,
67 deadline: Duration,
68}
69
70impl RequestScope {
71 #[must_use]
73 pub fn child(parent: &Cancellation, deadline: Duration) -> Self {
74 Self {
75 cancellation: parent.child(),
76 deadline,
77 }
78 }
79 #[must_use]
81 pub fn cancellation(&self) -> &Cancellation {
82 &self.cancellation
83 }
84 #[must_use]
86 pub fn deadline(&self) -> Duration {
87 self.deadline
88 }
89 pub fn cancel_timeout(&self) {
91 self.cancel("request deadline reached");
92 }
93 pub fn cancel_peer_drop(&self) {
95 self.cancel("peer disconnected");
96 }
97 fn cancel(&self, reason: &'static str) {
98 self.cancellation
99 .cancel(CancellationReason::new(reason).expect("static reason is valid"));
100 }
101}
102
103pub trait BodyReader {
105 fn next_chunk(&mut self, scope: &RequestScope) -> io::Result<Option<Vec<u8>>>;
107}
108
109pub trait ResponseWriter {
111 fn write_head(&mut self, head: ResponseHead, scope: &RequestScope) -> io::Result<()>;
113 fn write_chunk(&mut self, chunk: &[u8], scope: &RequestScope) -> io::Result<()>;
115 fn finish(&mut self, trailers: &[Header], scope: &RequestScope) -> io::Result<()>;
117}
118
119pub trait RawConnection {
121 fn parts(&mut self) -> (&RequestHead, &mut dyn BodyReader, &mut dyn ResponseWriter);
123}
124
125pub trait RawHandler: Send + Sync {
127 fn handle(
129 &self,
130 head: &RequestHead,
131 body: &mut dyn BodyReader,
132 response: &mut dyn ResponseWriter,
133 scope: &RequestScope,
134 ) -> Result<()>;
135}
136
137pub struct RawHttpServer<H> {
139 handler: H,
140 limits: BodyLimits,
141 trailers: TrailersPolicy,
142 request_deadline: Duration,
143 shutdown: Cancellation,
144 active: Mutex<Vec<Cancellation>>,
145}
146
147impl<H: RawHandler> RawHttpServer<H> {
148 pub fn new(
150 handler: H,
151 limits: BodyLimits,
152 trailers: TrailersPolicy,
153 request_deadline: Duration,
154 ) -> Result<Self> {
155 if request_deadline.is_zero() {
156 return Err(Error::Eval(
157 "raw HTTP request deadline must be non-zero".into(),
158 ));
159 }
160 Ok(Self {
161 handler,
162 limits: limits.validate()?,
163 trailers,
164 request_deadline,
165 shutdown: Cancellation::new(),
166 active: Mutex::new(Vec::new()),
167 })
168 }
169 pub fn shutdown(&self) {
171 self.shutdown
172 .cancel(CancellationReason::new("server shutdown").expect("static reason is valid"));
173 for request in self
174 .active
175 .lock()
176 .expect("active request mutex poisoned")
177 .drain(..)
178 {
179 request.cancel(
180 CancellationReason::new("server shutdown").expect("static reason is valid"),
181 );
182 }
183 }
184 pub fn serve(&self, connection: &mut dyn RawConnection, caller: &Cancellation) -> Result<()> {
186 let scope = RequestScope::child(caller, self.request_deadline);
187 if self.shutdown.is_cancelled() {
188 scope.cancel("server shutdown");
189 }
190 self.active
191 .lock()
192 .expect("active request mutex poisoned")
193 .push(scope.cancellation.clone());
194 let (head, body, response) = connection.parts();
195 let head = head.clone();
196 let mut body = LimitedBody {
197 inner: body,
198 limits: self.limits,
199 received: 0,
200 };
201 let body: &mut dyn BodyReader = &mut body;
202 let mut response = LimitedResponse {
203 inner: response,
204 max_chunk: self.limits.max_chunk_bytes,
205 trailers: self.trailers,
206 };
207 let result = self.handler.handle(&head, body, &mut response, &scope);
208 if result.is_err() {
209 scope.cancel("handler failure");
210 }
211 scope.cancel("request complete");
212 self.active
213 .lock()
214 .expect("active request mutex poisoned")
215 .retain(|request| !request.is_cancelled());
216 result
217 }
218}
219
220struct LimitedBody<'a> {
221 inner: &'a mut dyn BodyReader,
222 limits: BodyLimits,
223 received: usize,
224}
225impl BodyReader for LimitedBody<'_> {
226 fn next_chunk(&mut self, scope: &RequestScope) -> io::Result<Option<Vec<u8>>> {
227 if scope.cancellation().is_cancelled() {
228 return Err(io::Error::new(
229 io::ErrorKind::Interrupted,
230 "request cancelled",
231 ));
232 }
233 let chunk = self
234 .inner
235 .next_chunk(scope)
236 .inspect_err(|_| scope.cancel_peer_drop())?;
237 if let Some(chunk) = &chunk {
238 if chunk.is_empty()
239 || chunk.len() > self.limits.max_chunk_bytes
240 || self.received.saturating_add(chunk.len()) > self.limits.max_request_bytes
241 {
242 scope.cancel("request body cap exceeded");
243 return Err(io::Error::new(
244 io::ErrorKind::InvalidData,
245 "request body cap exceeded",
246 ));
247 }
248 self.received += chunk.len();
249 }
250 Ok(chunk)
251 }
252}
253
254struct LimitedResponse<'a> {
255 inner: &'a mut dyn ResponseWriter,
256 max_chunk: usize,
257 trailers: TrailersPolicy,
258}
259impl ResponseWriter for LimitedResponse<'_> {
260 fn write_head(&mut self, head: ResponseHead, scope: &RequestScope) -> io::Result<()> {
261 self.inner
262 .write_head(head, scope)
263 .inspect_err(|_| scope.cancel("response write failure"))
264 }
265 fn write_chunk(&mut self, chunk: &[u8], scope: &RequestScope) -> io::Result<()> {
266 if chunk.is_empty() || chunk.len() > self.max_chunk {
267 return Err(io::Error::new(
268 io::ErrorKind::InvalidInput,
269 "response chunk outside bounds",
270 ));
271 }
272 self.inner
273 .write_chunk(chunk, scope)
274 .inspect_err(|_| scope.cancel("response write failure"))
275 }
276 fn finish(&mut self, trailers: &[Header], scope: &RequestScope) -> io::Result<()> {
277 if !trailers.is_empty() && self.trailers == TrailersPolicy::Deny {
278 return Err(io::Error::new(
279 io::ErrorKind::InvalidInput,
280 "response trailers denied",
281 ));
282 }
283 self.inner
284 .finish(trailers, scope)
285 .inspect_err(|_| scope.cancel("response write failure"))
286 }
287}
288
289#[cfg(test)]
290mod tests {
291 use super::*;
292 use std::sync::{Arc, Mutex};
293
294 struct Body(Vec<Vec<u8>>);
295 impl BodyReader for Body {
296 fn next_chunk(&mut self, _: &RequestScope) -> io::Result<Option<Vec<u8>>> {
297 Ok(if self.0.is_empty() {
298 None
299 } else {
300 Some(self.0.remove(0))
301 })
302 }
303 }
304 #[derive(Default)]
305 struct Writer {
306 chunks: Vec<Vec<u8>>,
307 fail_after: usize,
308 }
309 impl ResponseWriter for Writer {
310 fn write_head(&mut self, _: ResponseHead, _: &RequestScope) -> io::Result<()> {
311 Ok(())
312 }
313 fn write_chunk(&mut self, chunk: &[u8], _: &RequestScope) -> io::Result<()> {
314 if self.chunks.len() == self.fail_after {
315 return Err(io::Error::new(io::ErrorKind::BrokenPipe, "peer dropped"));
316 }
317 self.chunks.push(chunk.to_vec());
318 Ok(())
319 }
320 fn finish(&mut self, _: &[Header], _: &RequestScope) -> io::Result<()> {
321 Ok(())
322 }
323 }
324 struct Connection {
325 head: RequestHead,
326 body: Body,
327 writer: Writer,
328 }
329 impl RawConnection for Connection {
330 fn parts(&mut self) -> (&RequestHead, &mut dyn BodyReader, &mut dyn ResponseWriter) {
331 (&self.head, &mut self.body, &mut self.writer)
332 }
333 }
334 struct Streaming {
335 observed: Arc<Mutex<Option<Cancellation>>>,
336 }
337 impl RawHandler for Streaming {
338 fn handle(
339 &self,
340 _: &RequestHead,
341 body: &mut dyn BodyReader,
342 out: &mut dyn ResponseWriter,
343 scope: &RequestScope,
344 ) -> Result<()> {
345 *self.observed.lock().unwrap() = Some(scope.cancellation().clone());
346 while let Some(chunk) = body
347 .next_chunk(scope)
348 .map_err(|e| Error::HostError(e.to_string()))?
349 {
350 out.write_chunk(&chunk, scope)
351 .map_err(|e| Error::HostError(e.to_string()))?;
352 }
353 Ok(())
354 }
355 }
356 #[test]
357 fn streaming_handler_is_backpressured_and_cancelled_on_peer_drop() {
358 let observed = Arc::new(Mutex::new(None));
359 let server = RawHttpServer::new(
360 Streaming {
361 observed: Arc::clone(&observed),
362 },
363 BodyLimits {
364 max_request_bytes: 16,
365 max_chunk_bytes: 4,
366 },
367 TrailersPolicy::Deny,
368 Duration::from_secs(1),
369 )
370 .unwrap();
371 let mut connection = Connection {
372 head: RequestHead {
373 method: "POST".into(),
374 target: "/mcp".into(),
375 headers: vec![("X-A".into(), "1".into()), ("X-A".into(), "2".into())],
376 peer: Some("peer".into()),
377 local: Some("local".into()),
378 },
379 body: Body(vec![b"one".to_vec(), b"two".to_vec()]),
380 writer: Writer {
381 fail_after: 1,
382 ..Writer::default()
383 },
384 };
385 assert!(server.serve(&mut connection, &Cancellation::new()).is_err());
386 assert_eq!(connection.writer.chunks, vec![b"one".to_vec()]);
387 assert!(observed.lock().unwrap().as_ref().unwrap().is_cancelled());
388 }
389}