1use std::any::{Any, TypeId};
2use std::collections::BTreeMap;
3use std::future::Future;
4use std::ops::Deref;
5use std::pin::Pin;
6use std::sync::{Arc, Weak};
7
8use unb_core::{Envelope, ErrorCode};
9use futures_util::{Stream, StreamExt};
10use serde::de::DeserializeOwned;
11use serde::Serialize;
12use serde_json::Value;
13
14use crate::handler::HandlerError;
15use crate::layer::{ErasedCall, Origin, ServiceBody};
16use crate::node::{Node, NodeSnapshot};
17
18#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
19pub enum Operation {
20 Unary,
21 Streaming,
22}
23
24impl Operation {
25 pub fn of<T>(request: &http::Request<T>) -> Option<Operation> {
26 let kind = request
27 .headers()
28 .get(unb_core::UNB_KIND)
29 .map(|value| value.to_str().ok());
30 match kind {
31 None | Some(Some("request")) | Some(Some("discover")) => Some(Operation::Unary),
32 Some(Some("subscribe")) | Some(Some("channel")) => Some(Operation::Streaming),
33 _ => None,
34 }
35 }
36}
37
38pub struct Request<T> {
39 payload: T,
40 parts: http::request::Parts,
41 origin: Origin,
42 node: Weak<Node>,
43 _snapshot: Arc<NodeSnapshot>,
44}
45
46impl<T> Request<T> {
47 pub fn payload(&self) -> &T {
48 &self.payload
49 }
50
51 pub fn into_payload(self) -> T {
52 self.payload
53 }
54
55 pub fn subject(&self) -> String {
56 Envelope::subject_of(&self.parts.uri)
57 }
58
59 pub fn method(&self) -> &http::Method {
60 &self.parts.method
61 }
62
63 pub fn headers(&self) -> &http::HeaderMap {
64 &self.parts.headers
65 }
66
67 pub fn extensions(&self) -> &http::Extensions {
68 &self.parts.extensions
69 }
70
71 pub fn take_body_stream(&self) -> Option<unb_runtime::BodyStream> {
72 self.parts
73 .extensions
74 .get::<StreamingBody>()?
75 .0
76 .lock()
77 .expect("streaming body slot")
78 .take()
79 }
80
81 pub fn origin(&self) -> &Origin {
82 &self.origin
83 }
84
85 pub async fn call(&self, subject: &str, payload: Value) -> Result<Value, HandlerError> {
86 let Some(node) = self.node.upgrade() else {
87 return Err(HandlerError::new(
88 ErrorCode::Internal,
89 "the node behind this request has shut down",
90 ));
91 };
92 let mut headers = serde_json::Map::new();
93 for (name, value) in &self.parts.headers {
94 if name.as_str().starts_with("unb-") {
95 continue;
96 }
97 let Ok(value) = std::str::from_utf8(value.as_bytes()) else {
98 continue;
99 };
100 headers.insert(name.as_str().to_string(), Value::String(value.to_string()));
101 }
102 node.call_nested(subject, payload, headers).await
103 }
104}
105
106impl<T: DeserializeOwned> Request<T> {
107 pub(crate) fn decode(request: http::Request<bytes::Bytes>) -> Result<Request<T>, HandlerError> {
108 let (parts, body) = request.into_parts();
109 let decoded = if body.is_empty() {
110 serde_json::from_value(Value::Null)
111 } else {
112 serde_json::from_slice(&body)
113 };
114 let payload: T = decoded.map_err(|error| {
115 let input = std::any::type_name::<T>()
116 .rsplit("::")
117 .next()
118 .unwrap_or("input");
119 let subject = Envelope::subject_of(&parts.uri);
120 HandlerError::new(
121 ErrorCode::InvalidInput,
122 format!(
123 "payload does not match {input}, the declared input for {subject:?}: {error}"
124 ),
125 )
126 })?;
127 let origin = parts
128 .extensions
129 .get::<Origin>()
130 .cloned()
131 .unwrap_or(Origin::Local);
132 let node = parts
133 .extensions
134 .get::<Weak<Node>>()
135 .cloned()
136 .unwrap_or_default();
137 let snapshot = parts
138 .extensions
139 .get::<Arc<NodeSnapshot>>()
140 .cloned()
141 .ok_or_else(|| HandlerError::new(ErrorCode::Internal, "missing invocation snapshot"))?;
142 Ok(Request {
143 payload,
144 parts,
145 origin,
146 node,
147 _snapshot: snapshot,
148 })
149 }
150}
151
152#[derive(Clone)]
153pub(crate) struct StreamingBody(
154 pub(crate) Arc<std::sync::Mutex<Option<unb_runtime::BodyStream>>>,
155);
156
157pub struct Reply<T>(T);
158
159impl<T> Reply<T> {
160 pub fn new(value: T) -> Reply<T> {
161 Reply(value)
162 }
163}
164
165enum StreamBody<T, E> {
166 Typed(Pin<Box<dyn Stream<Item = Result<T, E>> + Send>>),
167 Raw(crate::handler::EventStream),
168}
169
170pub struct Streaming<T, E> {
171 body: StreamBody<T, E>,
172}
173
174impl<T, E> Streaming<T, E> {
175 pub fn new(stream: impl Stream<Item = Result<T, E>> + Send + 'static) -> Streaming<T, E> {
176 Streaming {
177 body: StreamBody::Typed(Box::pin(stream)),
178 }
179 }
180
181 pub fn raw(
182 stream: impl Stream<Item = Result<bytes::Bytes, HandlerError>> + Send + 'static,
183 ) -> Streaming<T, E> {
184 Streaming {
185 body: StreamBody::Raw(Box::pin(stream)),
186 }
187 }
188}
189
190mod sealed {
191 pub trait Sealed {}
192}
193
194pub trait HandlerOutput: sealed::Sealed {
195 const OPERATION: Operation;
196 fn into_response(self) -> Result<http::Response<ServiceBody>, HandlerError>;
197}
198
199fn respond(body: ServiceBody) -> Result<http::Response<ServiceBody>, HandlerError> {
200 http::Response::builder().body(body).map_err(|error| {
201 HandlerError::new(
202 ErrorCode::Internal,
203 format!("response construction failed: {error}"),
204 )
205 })
206}
207
208impl<T: Serialize> sealed::Sealed for Reply<T> {}
209
210impl<T: Serialize> HandlerOutput for Reply<T> {
211 const OPERATION: Operation = Operation::Unary;
212
213 fn into_response(self) -> Result<http::Response<ServiceBody>, HandlerError> {
214 let value = serde_json::to_value(self.0).map_err(|error| {
215 HandlerError::new(
216 ErrorCode::Internal,
217 format!("response serialization failed: {error}"),
218 )
219 })?;
220 respond(ServiceBody::Unary(Envelope::encode_payload(&value)))
221 }
222}
223
224impl<T, E> sealed::Sealed for Streaming<T, E>
225where
226 T: Serialize + Send + 'static,
227 E: Into<HandlerError> + Send + 'static,
228{
229}
230
231impl<T, E> HandlerOutput for Streaming<T, E>
232where
233 T: Serialize + Send + 'static,
234 E: Into<HandlerError> + Send + 'static,
235{
236 const OPERATION: Operation = Operation::Streaming;
237
238 fn into_response(self) -> Result<http::Response<ServiceBody>, HandlerError> {
239 let body = match self.body {
240 StreamBody::Typed(stream) => ServiceBody::Stream(Box::pin(stream.map(|item| {
241 match item {
242 Ok(event) => serde_json::to_value(event)
243 .map(|value| Envelope::encode_payload(&value))
244 .map_err(|error| {
245 HandlerError::new(
246 ErrorCode::Internal,
247 format!("event serialization failed: {error}"),
248 )
249 }),
250 Err(error) => Err(error.into()),
251 }
252 }))),
253 StreamBody::Raw(stream) => ServiceBody::Stream(stream),
254 };
255 respond(body)
256 }
257}
258
259pub struct State<T>(Arc<T>);
260
261impl<T> State<T> {
262 pub fn new(value: T) -> State<T> {
263 State(Arc::new(value))
264 }
265}
266
267impl<T> Clone for State<T> {
268 fn clone(&self) -> State<T> {
269 State(self.0.clone())
270 }
271}
272
273impl<T> Deref for State<T> {
274 type Target = T;
275
276 fn deref(&self) -> &T {
277 &self.0
278 }
279}
280
281#[derive(Default, Clone)]
282pub(crate) struct StateMap {
283 values: BTreeMap<TypeId, Arc<dyn Any + Send + Sync>>,
284}
285
286impl StateMap {
287 pub(crate) fn insert<T: Send + Sync + 'static>(&mut self, value: T) {
288 self.values.insert(TypeId::of::<T>(), Arc::new(value));
289 }
290
291 pub(crate) fn get<T: Send + Sync + 'static>(&self) -> Option<State<T>> {
292 self.values
293 .get(&TypeId::of::<T>())
294 .cloned()
295 .and_then(|any| any.downcast::<T>().ok())
296 .map(State)
297 }
298
299 pub(crate) fn merged_over(&self, outer: &StateMap) -> StateMap {
300 let mut merged = outer.clone();
301 for (key, value) in &self.values {
302 merged.values.insert(*key, value.clone());
303 }
304 merged
305 }
306}
307
308pub struct States<'a>(pub(crate) &'a StateMap);
309
310impl States<'_> {
311 pub fn state<T: Send + Sync + 'static>(&self) -> Result<State<T>, String> {
312 self.0.get::<T>().ok_or_else(|| {
313 format!(
314 "no registered state provides {}; register it with state(...) on the node or scope",
315 std::any::type_name::<T>()
316 )
317 })
318 }
319}
320
321pub enum ContractSchema {
322 Static(fn() -> Value),
323 Owned(Value),
324}
325
326impl From<fn() -> Value> for ContractSchema {
327 fn from(factory: fn() -> Value) -> ContractSchema {
328 ContractSchema::Static(factory)
329 }
330}
331
332impl From<Value> for ContractSchema {
333 fn from(value: Value) -> ContractSchema {
334 ContractSchema::Owned(value)
335 }
336}
337
338impl ContractSchema {
339 fn value(&self) -> Value {
340 match self {
341 ContractSchema::Static(factory) => factory(),
342 ContractSchema::Owned(value) => value.clone(),
343 }
344 }
345}
346
347pub struct OperationContract {
348 pub input: Option<ContractSchema>,
349 pub output: Option<ContractSchema>,
350 pub event: Option<ContractSchema>,
351 pub error: Option<ContractSchema>,
352}
353
354impl OperationContract {
355 pub fn unknown() -> OperationContract {
356 OperationContract {
357 input: None,
358 output: None,
359 event: None,
360 error: None,
361 }
362 }
363
364 pub(crate) fn to_json(&self, operation: Operation) -> Value {
365 let render = |schema: &Option<ContractSchema>| {
366 schema
367 .as_ref()
368 .map(ContractSchema::value)
369 .unwrap_or_else(|| serde_json::json!({ "unknown": true }))
370 };
371 match operation {
372 Operation::Unary => serde_json::json!({
373 "input_schema": render(&self.input),
374 "output_schema": render(&self.output),
375 }),
376 Operation::Streaming => serde_json::json!({
377 "input_schema": render(&self.input),
378 "event_schema": render(&self.event),
379 "error_schema": render(&self.error),
380 }),
381 }
382 }
383}
384
385type BuildFn = Box<dyn FnOnce(&States<'_>) -> Result<ErasedCall, String> + Send>;
386
387pub struct HandlerService {
388 pub(crate) local_name: String,
389 pub(crate) subject_override: Option<String>,
390 pub(crate) one_line: Option<String>,
391 pub(crate) operation: Operation,
392 pub(crate) metadata: Option<Value>,
393 pub(crate) contract: OperationContract,
394 pub(crate) build: BuildFn,
395}
396
397impl HandlerService {
398 pub fn declare(
399 local_name: &str,
400 subject_override: Option<&str>,
401 one_line: Option<&str>,
402 operation: Operation,
403 contract: OperationContract,
404 build: impl FnOnce(&States<'_>) -> Result<ErasedCall, String> + Send + 'static,
405 ) -> HandlerService {
406 HandlerService {
407 local_name: local_name.into(),
408 subject_override: subject_override.map(Into::into),
409 one_line: one_line.map(Into::into),
410 operation,
411 metadata: None,
412 contract,
413 build: Box::new(build),
414 }
415 }
416
417 pub fn at_subject(mut self, subject: impl Into<String>) -> HandlerService {
418 self.subject_override = Some(subject.into());
419 self
420 }
421
422 pub fn describe(mut self, metadata: Value) -> HandlerService {
423 self.metadata = Some(metadata);
424 self
425 }
426
427 pub(crate) fn effective_subject(&self, scopes: &[String]) -> Result<String, String> {
428 let local = self.subject_override.as_deref().unwrap_or(&self.local_name);
429 let subject = if scopes.is_empty() {
430 local.to_string()
431 } else {
432 format!("{}.{local}", scopes.join("."))
433 };
434 let valid = !subject.is_empty()
435 && subject.len() <= unb_core::MAX_SUBJECT_LEN
436 && subject.split('.').all(|segment| !segment.is_empty());
437 if valid {
438 Ok(subject)
439 } else {
440 Err(format!(
441 "subject {subject:?} needs non-empty dot-separated segments within {} bytes",
442 unb_core::MAX_SUBJECT_LEN
443 ))
444 }
445 }
446}
447
448pub trait Handler: Sized {
449 fn into_service(self) -> HandlerService;
450
451 fn at_subject(self, subject: impl Into<String>) -> HandlerService {
452 self.into_service().at_subject(subject)
453 }
454
455 fn describe(self, metadata: Value) -> HandlerService {
456 self.into_service().describe(metadata)
457 }
458}
459
460impl Handler for HandlerService {
461 fn into_service(self) -> HandlerService {
462 self
463 }
464}
465
466pub fn erase_unary<F, Fut, In, Out, E>(f: F) -> ErasedCall
467where
468 F: Fn(Request<In>) -> Fut + Send + Sync + 'static,
469 Fut: Future<Output = Result<Reply<Out>, E>> + Send + 'static,
470 In: DeserializeOwned + Send + 'static,
471 Out: Serialize + Send + 'static,
472 E: Into<HandlerError> + Send + 'static,
473{
474 let f = Arc::new(f);
475 Arc::new(move |request: http::Request<bytes::Bytes>| {
476 let f = f.clone();
477 Box::pin(async move {
478 let request = Request::<In>::decode(request)?;
479 match f(request).await {
480 Ok(reply) => reply.into_response(),
481 Err(error) => Err(error.into()),
482 }
483 })
484 })
485}
486
487pub fn erase_streaming<F, Fut, In, Event, StreamError, E>(f: F) -> ErasedCall
488where
489 F: Fn(Request<In>) -> Fut + Send + Sync + 'static,
490 Fut: Future<Output = Result<Streaming<Event, StreamError>, E>> + Send + 'static,
491 In: DeserializeOwned + Send + 'static,
492 Event: Serialize + Send + 'static,
493 StreamError: Into<HandlerError> + Send + 'static,
494 E: Into<HandlerError> + Send + 'static,
495{
496 let f = Arc::new(f);
497 Arc::new(move |request: http::Request<bytes::Bytes>| {
498 let f = f.clone();
499 Box::pin(async move {
500 let request = Request::<In>::decode(request)?;
501 match f(request).await {
502 Ok(streaming) => streaming.into_response(),
503 Err(error) => Err(error.into()),
504 }
505 })
506 })
507}
508
509#[cfg(test)]
510mod tests {
511 use serde::Deserialize;
512 use serde_json::json;
513
514 use super::*;
515
516 fn service_request(payload: Value) -> http::Request<bytes::Bytes> {
517 let node = Node::builder("service-test")
518 .insecure_accept_declared_peer_identities()
519 .build()
520 .expect("test node builds");
521 let mut request = http::Request::builder()
522 .method("POST")
523 .uri("/probe")
524 .body(Envelope::encode_payload(&payload))
525 .expect("test request is well formed");
526 request.extensions_mut().insert(Origin::Local);
527 request.extensions_mut().insert(node.snapshot.load_full());
528 request
529 }
530
531 #[derive(Deserialize)]
532 struct Input {
533 n: u32,
534 }
535
536 #[derive(Serialize)]
537 struct Output {
538 doubled: u32,
539 }
540
541 #[tokio::test]
542 async fn a_unary_reply_serializes_through_the_erased_adapter() {
543 let call = erase_unary(|request: Request<Input>| async move {
544 Ok::<_, HandlerError>(Reply::new(Output {
545 doubled: request.payload().n * 2,
546 }))
547 });
548 let response = call(service_request(json!({ "n": 21 })))
549 .await
550 .unwrap_or_else(|error| panic!("call failed: {error}"));
551 let ServiceBody::Unary(payload) = response.into_body() else {
552 panic!("expected a unary response");
553 };
554 let value: Value = serde_json::from_slice(&payload).expect("unary payload is json");
555 assert_eq!(value["doubled"], 42);
556 }
557
558 #[tokio::test]
559 async fn an_undeclared_payload_fails_decode_before_the_handler_runs() {
560 let call = erase_unary(|_request: Request<Input>| async move {
561 panic!("the handler must not run on an invalid payload");
562 #[allow(unreachable_code)]
563 Ok::<_, HandlerError>(Reply::new(Value::Null))
564 });
565 let error = match call(service_request(json!({ "n": "not-a-number" }))).await {
566 Err(error) => error,
567 Ok(_) => panic!("decode must fail"),
568 };
569 assert_eq!(error.code, ErrorCode::InvalidInput);
570 assert!(error.message.contains("probe"));
571 }
572
573 #[tokio::test]
574 async fn a_streaming_output_maps_events_and_errors_into_the_event_stream() {
575 let call = erase_streaming(|_request: Request<Input>| async move {
576 let events = futures_util::stream::iter(vec![
577 Ok(Output { doubled: 2 }),
578 Err(HandlerError::new(ErrorCode::Internal, "stream broke")),
579 ]);
580 Ok::<_, HandlerError>(Streaming::new(events))
581 });
582 let response = call(service_request(json!({ "n": 1 })))
583 .await
584 .unwrap_or_else(|error| panic!("call failed: {error}"));
585 let ServiceBody::Stream(mut stream) = response.into_body() else {
586 panic!("expected a stream response");
587 };
588 let first = stream.next().await.unwrap().unwrap();
589 let first: Value = serde_json::from_slice(&first).unwrap();
590 assert_eq!(first["doubled"], 2);
591 let second = stream.next().await.unwrap().unwrap_err();
592 assert_eq!(second.code, ErrorCode::Internal);
593 assert!(stream.next().await.is_none());
594 }
595
596 #[test]
597 fn nearest_scope_state_wins_over_outer_state() {
598 let mut node_states = StateMap::default();
599 node_states.insert(7u32);
600 node_states.insert("node".to_string());
601 let mut scope_states = StateMap::default();
602 scope_states.insert("scope".to_string());
603 let merged = scope_states.merged_over(&node_states);
604 assert_eq!(*merged.get::<String>().unwrap(), "scope");
605 assert_eq!(*merged.get::<u32>().unwrap(), 7);
606 let missing = match States(&merged).state::<bool>() {
607 Err(missing) => missing,
608 Ok(_) => panic!("bool state must be absent"),
609 };
610 assert!(missing.contains("bool"));
611 }
612}