icap_rs/server/builder.rs
1use std::collections::HashMap;
2use std::collections::hash_map::Entry;
3use std::future::Future;
4use std::sync::Arc;
5
6use tokio::net::TcpListener;
7use tokio::sync::Semaphore;
8
9use tokio_util::task::TaskTracker;
10
11use crate::error::IcapResult;
12use crate::request::{IncomingRequest, RequestParserMode, normalize_service_path};
13use crate::server::timeouts::ServerTimeouts;
14use crate::server::{ShutdownEvent, default_shutdown_handler};
15#[cfg(feature = "tls-rustls")]
16use crate::tls::ServerTlsConfig;
17use crate::{Method, ServiceOptions};
18
19use super::Server;
20use super::router::{HandlerEntry, RequestHandler, RouteEntry, RouteOutput, resolve_service};
21
22/// Builder for [`Server`].
23///
24/// Construct with [`Server::builder`], register routes, then call [`ServerBuilder::build`].
25///
26/// Duplicate registration of the **same (service, method)** panics with a clear error,
27/// like axum.
28#[must_use]
29pub struct ServerBuilder {
30 bind_addr: Option<String>,
31 routes: HashMap<String, RouteEntry>,
32 max_connections_global: Option<usize>,
33 aliases: HashMap<String, String>,
34 default_service: Option<String>,
35 request_parser_mode: RequestParserMode,
36 timeouts: ServerTimeouts,
37 max_request_header_bytes: Option<usize>,
38 shutdown_handler: Arc<dyn Fn(ShutdownEvent) + Send + Sync>,
39 task_tracker: Option<TaskTracker>,
40 #[cfg(feature = "tls-rustls")]
41 tls: Option<ServerTlsConfig>,
42}
43
44impl Default for ServerBuilder {
45 fn default() -> Self {
46 Self {
47 bind_addr: None,
48 routes: HashMap::new(),
49 max_connections_global: None,
50 aliases: HashMap::new(),
51 default_service: None,
52 request_parser_mode: RequestParserMode::default(),
53 timeouts: ServerTimeouts::default(),
54 max_request_header_bytes: None,
55 shutdown_handler: Arc::new(default_shutdown_handler),
56 task_tracker: None,
57 #[cfg(feature = "tls-rustls")]
58 tls: None,
59 }
60 }
61}
62
63impl ServerBuilder {
64 /// Enable ICAPS (ICAP over TLS) using the supplied [`ServerTlsConfig`].
65 ///
66 /// Configuration objects encapsulate the certificate chain, private key,
67 /// optional client-certificate verification and the handshake timeout.
68 /// See the [`crate::tls`] module for builders to construct one.
69 ///
70 /// Only available when the `tls-rustls` feature is enabled.
71 ///
72 /// # Example
73 /// ```no_run
74 /// use icap_rs::{
75 /// IcapResult, IncomingRequest, Method, Response, Server, ServerTlsConfig, ServiceOptions,
76 /// };
77 ///
78 /// const ISTAG: &str = "scan-1.0";
79 ///
80 /// #[tokio::main]
81 /// async fn main() -> IcapResult<()> {
82 /// let tls = ServerTlsConfig::from_pem_files("certs/server.crt", "certs/server.key")?;
83 ///
84 /// let server = Server::builder()
85 /// .bind("0.0.0.0:11344")
86 /// .with_tls(tls)
87 /// .route(
88 /// "scan",
89 /// [Method::ReqMod, Method::RespMod],
90 /// |_req: IncomingRequest| async move {
91 /// Ok(Response::no_content_with_istag(ISTAG)?)
92 /// },
93 /// Some(ServiceOptions::new()
94 /// .with_static_istag(ISTAG)
95 /// .with_preview(2048)
96 /// .allow_204()),
97 /// )
98 /// .build().await?;
99 ///
100 /// server.run().await
101 /// }
102 /// ```
103 #[cfg(feature = "tls-rustls")]
104 pub fn with_tls(mut self, config: ServerTlsConfig) -> Self {
105 self.tls = Some(config);
106 self
107 }
108
109 /// Set the bind address, e.g. `"127.0.0.1:1344"`.
110 pub fn bind(mut self, addr: &str) -> Self {
111 self.bind_addr = Some(addr.to_string());
112 self
113 }
114
115 /// Limit the number of concurrent connections accepted by the server.
116 ///
117 /// The value is also advertised in `OPTIONS` as `Max-Connections` if
118 /// not explicitly configured on the per-service options.
119 pub fn with_max_connections(mut self, n: usize) -> Self {
120 self.max_connections_global = Some(n.max(1));
121 self
122 }
123
124 /// Install a full [`ServerTimeouts`] configuration in one shot.
125 ///
126 /// See [`ServerTimeouts`] for the meaning of each field. Defaults are
127 /// `None` (no timeout); fields not set on the supplied value disable the
128 /// corresponding deadline.
129 pub const fn with_timeouts(mut self, timeouts: ServerTimeouts) -> Self {
130 self.timeouts = timeouts;
131 self
132 }
133
134 /// Set the maximum ICAP request header block size, in bytes.
135 ///
136 /// The limit includes the request line, all ICAP header lines, and the
137 /// terminating `CRLFCRLF`. The default is 64 KiB. Oversized request headers
138 /// receive `400 Bad Request` and the connection is closed.
139 pub const fn with_request_header_limit(mut self, bytes: usize) -> Self {
140 self.max_request_header_bytes = Some(bytes);
141 self
142 }
143
144 /// Register a callback that is called with [`ShutdownEvent`] during graceful shutdown.
145 ///
146 /// The handler runs synchronously inside the accept loop task — keep it fast.
147 /// Use it for custom logging, metrics, or alerting. When not set, the server
148 /// logs via [`tracing::warn`] by default.
149 ///
150 /// # Example
151 ///
152 /// ```rust,no_run
153 /// use icap_rs::{IcapResult, Server, ShutdownEvent};
154 ///
155 /// #[tokio::main]
156 /// async fn main() -> IcapResult<()> {
157 /// let server = Server::builder()
158 /// .bind("127.0.0.1:1344")
159 /// .on_shutdown_event(|event| match event {
160 /// ShutdownEvent::Draining { active_connections, drain_timeout } => {
161 /// eprintln!("[shutdown] {active_connections} connection(s) still active");
162 /// if let Some(d) = drain_timeout {
163 /// eprintln!("[shutdown] force-close in {d:.1?}");
164 /// }
165 /// }
166 /// ShutdownEvent::DrainTimedOut { remaining_connections } => {
167 /// eprintln!("[shutdown] timed out, cancelling {remaining_connections}");
168 /// }
169 /// _ => {}
170 /// })
171 /// .build()
172 /// .await?;
173 ///
174 /// server.run_until(async { tokio::signal::ctrl_c().await.ok(); }).await
175 /// }
176 /// ```
177 pub fn on_shutdown_event<F>(mut self, handler: F) -> Self
178 where
179 F: Fn(ShutdownEvent) + Send + Sync + 'static,
180 {
181 self.shutdown_handler = Arc::new(handler);
182 self
183 }
184
185 /// Enable legacy compatibility request parsing.
186 ///
187 /// Strict RFC parsing is the default and requires every ICAP request,
188 /// including `OPTIONS`, to carry an `Encapsulated` header. This opt-in mode
189 /// accepts legacy `OPTIONS` requests without `Encapsulated`.
190 pub const fn with_compatibility_request_parser(mut self) -> Self {
191 self.request_parser_mode = RequestParserMode::Compatibility;
192 self
193 }
194
195 /// Register a **service route** for one or more ICAP methods.
196 ///
197 /// - Each service must have a [`ServiceOptions`] value with an explicit
198 /// `ISTag`; routes without options are rejected by [`build`](Self::build).
199 /// - Multiple calls to `.route(..)` for the **same service** are allowed as long as methods do not overlap.
200 /// - Registering the **same method** for the same service twice will `panic!` with a clear message.
201 /// - The same handler can be reused for multiple methods in a single call.
202 /// - Return `IcapResult<PreviewDecision>` from the handler to make the route
203 /// preview-aware. Such handlers are called with `Body::Preview` after
204 /// preview bytes arrive and before the server sends `100 Continue`.
205 /// Returning `PreviewDecision::Continue` resumes the RFC preview flow; the
206 /// same handler is called again with `Body::Full` after the remainder is read.
207 ///
208 /// ## Handler invocation and the `Allow: 204` header
209 ///
210 /// The handler is **always called** for every REQMOD/RESPMOD request.
211 ///
212 /// RFC 3507 §4.6 prohibits the server from returning `204 No Content` unless
213 /// the client explicitly advertised `Allow: 204`. When the handler returns
214 /// `204` but the client did not send `Allow: 204`, the server automatically
215 /// converts the response: it echoes the original embedded HTTP message in a
216 /// `200 OK` (or `206 Partial Content` if `Allow: 206` was sent).
217 ///
218 /// This means handlers can always return `Response::no_content()` to signal
219 /// "no modification needed" — the server takes care of the RFC-compliant
220 /// wrapping regardless of what the client advertised.
221 pub fn route<MIt, MItem, F, Fut>(
222 mut self,
223 service: &str,
224 methods: MIt,
225 handler: F,
226 options: Option<ServiceOptions>,
227 ) -> Self
228 where
229 MIt: IntoIterator<Item = MItem>,
230 MItem: Into<Method>,
231 F: Fn(IncomingRequest) -> Fut + Send + Sync + 'static,
232 Fut: Future + Send + 'static,
233 Fut::Output: RouteOutput,
234 {
235 // Routing is by full request path (RFC 3507 §6.4): `/v1/scan` and
236 // `/v2/scan` are distinct services. Normalize so the registered key
237 // matches the path the parser extracts from the request line.
238 let service = normalize_service_path(service);
239 let entry = match self.routes.entry(service.clone()) {
240 Entry::Occupied(o) => o.into_mut(),
241 Entry::Vacant(v) => v.insert(RouteEntry {
242 handlers: HashMap::new(),
243 options: None,
244 }),
245 };
246
247 // Wrap handler in Arc so we can reuse it for multiple methods
248 let h_arc = Arc::new(handler);
249
250 for item in methods {
251 let m: Method = item.into();
252
253 assert_ne!(
254 m,
255 Method::Options,
256 "OPTIONS cannot have a handler; it's answered automatically for '{service}'"
257 );
258 assert!(
259 !entry.handlers.contains_key(&m),
260 "Overlapping method route. Handler for '{m} {service}' already exists"
261 );
262
263 let h_clone = h_arc.clone();
264 let h: RequestHandler = Box::new(move |req| {
265 let h = h_clone.clone();
266 Box::pin(async move { h(req).await.into_preview_decision() })
267 });
268 entry.handlers.insert(
269 m,
270 HandlerEntry {
271 handler: h,
272 preview_aware: Fut::Output::PREVIEW_AWARE,
273 },
274 );
275 }
276
277 // Attach options if provided for this route
278 if let Some(cfg) = options {
279 assert!(
280 entry.options.is_none(),
281 "Options already set for service '{service}'"
282 );
283 entry.options = Some(cfg);
284 }
285
286 self
287 }
288
289 /// Register a route for `REQMOD` only.
290 ///
291 /// Convenience wrapper around [`route`](Self::route) with `methods = [Method::ReqMod]`.
292 /// See [`route`](Self::route) for full semantics, including panic conditions on
293 /// duplicate (service, method) registration and `ServiceOptions` requirements.
294 ///
295 /// # Panics
296 ///
297 /// Panics if a `REQMOD` handler for the same `service` was already registered,
298 /// or if `options` is provided more than once for the same service.
299 pub fn route_reqmod<F, Fut>(
300 self,
301 service: &str,
302 handler: F,
303 options: Option<ServiceOptions>,
304 ) -> Self
305 where
306 F: Fn(IncomingRequest) -> Fut + Send + Sync + 'static,
307 Fut: Future + Send + 'static,
308 Fut::Output: RouteOutput,
309 {
310 self.route(service, [Method::ReqMod], handler, options)
311 }
312
313 /// Register a route for `RESPMOD` only.
314 ///
315 /// Convenience wrapper around [`route`](Self::route) with `methods = [Method::RespMod]`.
316 /// See [`route`](Self::route) for full semantics, including panic conditions on
317 /// duplicate (service, method) registration and `ServiceOptions` requirements.
318 ///
319 /// # Panics
320 ///
321 /// Panics if a `RESPMOD` handler for the same `service` was already registered,
322 /// or if `options` is provided more than once for the same service.
323 pub fn route_respmod<F, Fut>(
324 self,
325 service: &str,
326 handler: F,
327 options: Option<ServiceOptions>,
328 ) -> Self
329 where
330 F: Fn(IncomingRequest) -> Fut + Send + Sync + 'static,
331 Fut: Future + Send + 'static,
332 Fut::Output: RouteOutput,
333 {
334 self.route(service, [Method::RespMod], handler, options)
335 }
336
337 /// Add an alias for a service path: `from` → `to`.
338 ///
339 /// Both ends are normalized to canonical request paths, so `"scan"` and
340 /// `"/scan"` refer to the same route. Useful to make the root path behave
341 /// like an existing service:
342 /// ```
343 /// # use icap_rs::Server;
344 /// let builder = Server::builder()
345 /// .alias("/", "scan"); // "icap://host:1344/" is routed to "/scan"
346 /// ```
347 ///
348 /// Notes:
349 /// - Aliases are applied *after* [`default_service`](Self::default_service) is considered
350 /// for the root ("/") path.
351 /// - Up to 4 alias rewrites are applied to avoid cycles.
352 pub fn alias(mut self, from: &str, to: &str) -> Self {
353 self.aliases
354 .insert(normalize_service_path(from), normalize_service_path(to));
355 self
356 }
357
358 /// Set a default service for the root ("/") path (e.g. `"scan"`).
359 ///
360 /// The value is normalized to a canonical request path, so `"scan"` and
361 /// `"/scan"` are equivalent.
362 ///
363 /// Example:
364 /// ```
365 /// # use icap_rs::Server;
366 /// let builder = Server::builder()
367 /// .default_service("scan");
368 /// ```
369 ///
370 /// If a client sends `icap://host:1344/` or an empty service, requests are internally
371 /// routed to the specified service path.
372 pub fn default_service(mut self, svc: &str) -> Self {
373 self.default_service = Some(normalize_service_path(svc));
374 self
375 }
376
377 /// Register a [`TaskTracker`] for user-owned background tasks.
378 ///
379 /// After all active connections drain following a shutdown signal, the server
380 /// calls [`TaskTracker::close`] on the tracker and waits for all tracked tasks
381 /// to finish before returning from [`Server::run_until`].
382 ///
383 /// If a drain timeout is configured via [`ServerTimeouts::with_shutdown_drain`],
384 /// the remaining budget is shared: once the drain deadline fires the tracker
385 /// wait is skipped and the server returns immediately.
386 ///
387 /// # Example
388 ///
389 /// ```rust,no_run
390 /// use icap_rs::{IcapResult, Server};
391 /// use tokio_util::task::TaskTracker;
392 ///
393 /// #[tokio::main]
394 /// async fn main() -> IcapResult<()> {
395 /// let tracker = TaskTracker::new();
396 ///
397 /// // Spawn a background task and track it so the server waits for it on shutdown.
398 /// tracker.spawn(async {
399 /// // background work ...
400 /// });
401 ///
402 /// let server = Server::builder()
403 /// .bind("127.0.0.1:1344")
404 /// .with_task_tracker(tracker)
405 /// .build()
406 /// .await?;
407 ///
408 /// server.run_until(async { tokio::signal::ctrl_c().await.ok(); }).await
409 /// }
410 /// ```
411 pub fn with_task_tracker(mut self, tracker: TaskTracker) -> Self {
412 self.task_tracker = Some(tracker);
413 self
414 }
415
416 /// Finalize the builder and create a [`Server`].
417 pub async fn build(self) -> IcapResult<Server> {
418 validate_builder_config(&self.routes, &self.aliases, self.default_service.as_deref())?;
419
420 let bind_addr = self
421 .bind_addr
422 .unwrap_or_else(|| "127.0.0.1:1344".to_string());
423 let listener = TcpListener::bind(&bind_addr).await?;
424
425 let conn_limit = self
426 .max_connections_global
427 .map(|n| Arc::new(Semaphore::new(n)));
428 let advertised_max_conn = self.max_connections_global;
429
430 // TLS (only when feature is enabled)
431 #[cfg(feature = "tls-rustls")]
432 let tls = self.tls.map(ServerTlsConfig::into_acceptor).transpose()?;
433
434 Ok(Server {
435 listener,
436 routes: Arc::new(self.routes),
437 conn_limit,
438 advertised_max_conn,
439 aliases: Arc::new(self.aliases),
440 default_service: self.default_service,
441 request_parser_mode: self.request_parser_mode,
442 timeouts: self.timeouts,
443 max_request_header_bytes: self
444 .max_request_header_bytes
445 .unwrap_or(crate::DEFAULT_ICAP_HEADER_BYTES),
446 shutdown_handler: self.shutdown_handler,
447 task_tracker: self.task_tracker,
448
449 #[cfg(feature = "tls-rustls")]
450 tls,
451 })
452 }
453}
454
455fn validate_builder_config(
456 routes: &HashMap<String, RouteEntry>,
457 aliases: &HashMap<String, String>,
458 default_service: Option<&str>,
459) -> IcapResult<()> {
460 use crate::error::ConfigError;
461
462 if let Some(default) = default_service {
463 let resolved = resolve_service(default, aliases, None);
464 if !routes.contains_key(resolved.as_ref()) {
465 return Err(ConfigError::UnknownDefaultService {
466 name: default.to_owned(),
467 resolved: resolved.into_owned(),
468 }
469 .into());
470 }
471 }
472
473 for (from, to) in aliases {
474 let resolved = resolve_service(to, aliases, None);
475 if !routes.contains_key(resolved.as_ref()) {
476 return Err(ConfigError::UnknownAlias {
477 from: from.clone(),
478 resolved: resolved.into_owned(),
479 }
480 .into());
481 }
482 }
483
484 for (service, entry) in routes {
485 if entry.handlers.is_empty() {
486 return Err(ConfigError::ServiceWithoutHandlers {
487 service: service.clone(),
488 }
489 .into());
490 }
491 let Some(options) = &entry.options else {
492 return Err(ConfigError::MissingServiceOptions {
493 service: service.clone(),
494 }
495 .into());
496 };
497 if let Err(err) = options.validate() {
498 return Err(ConfigError::InvalidServiceOptions {
499 service: service.clone(),
500 reason: err,
501 }
502 .into());
503 }
504 }
505
506 Ok(())
507}