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 futures_util::{Stream, StreamExt};
9use serde::de::DeserializeOwned;
10use serde::Serialize;
11use serde_json::Value;
12use unb_core::{Envelope, ErrorCode};
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(pub(crate) Arc<std::sync::Mutex<Option<unb_runtime::BodyStream>>>);
154
155pub struct Reply<T>(T);
156
157impl<T> Reply<T> {
158 pub fn new(value: T) -> Reply<T> {
159 Reply(value)
160 }
161}
162
163enum StreamBody<T, E> {
164 Typed(Pin<Box<dyn Stream<Item = Result<T, E>> + Send>>),
165 Raw(crate::handler::EventStream),
166}
167
168pub struct Streaming<T, E> {
169 body: StreamBody<T, E>,
170}
171
172impl<T, E> Streaming<T, E> {
173 pub fn new(stream: impl Stream<Item = Result<T, E>> + Send + 'static) -> Streaming<T, E> {
174 Streaming {
175 body: StreamBody::Typed(Box::pin(stream)),
176 }
177 }
178
179 pub fn raw(
180 stream: impl Stream<Item = Result<bytes::Bytes, HandlerError>> + Send + 'static,
181 ) -> Streaming<T, E> {
182 Streaming {
183 body: StreamBody::Raw(Box::pin(stream)),
184 }
185 }
186}
187
188mod sealed {
189 pub trait Sealed {}
190}
191
192pub trait HandlerOutput: sealed::Sealed {
193 const OPERATION: Operation;
194 fn into_response(self) -> Result<http::Response<ServiceBody>, HandlerError>;
195}
196
197fn respond(body: ServiceBody) -> Result<http::Response<ServiceBody>, HandlerError> {
198 http::Response::builder().body(body).map_err(|error| {
199 HandlerError::new(
200 ErrorCode::Internal,
201 format!("response construction failed: {error}"),
202 )
203 })
204}
205
206impl<T: Serialize> sealed::Sealed for Reply<T> {}
207
208impl<T: Serialize> HandlerOutput for Reply<T> {
209 const OPERATION: Operation = Operation::Unary;
210
211 fn into_response(self) -> Result<http::Response<ServiceBody>, HandlerError> {
212 let value = serde_json::to_value(self.0).map_err(|error| {
213 HandlerError::new(
214 ErrorCode::Internal,
215 format!("response serialization failed: {error}"),
216 )
217 })?;
218 respond(ServiceBody::Unary(Envelope::encode_payload(&value)))
219 }
220}
221
222impl<T, E> sealed::Sealed for Streaming<T, E>
223where
224 T: Serialize + Send + 'static,
225 E: Into<HandlerError> + Send + 'static,
226{
227}
228
229impl<T, E> HandlerOutput for Streaming<T, E>
230where
231 T: Serialize + Send + 'static,
232 E: Into<HandlerError> + Send + 'static,
233{
234 const OPERATION: Operation = Operation::Streaming;
235
236 fn into_response(self) -> Result<http::Response<ServiceBody>, HandlerError> {
237 let body = match self.body {
238 StreamBody::Typed(stream) => ServiceBody::Stream(Box::pin(stream.map(|item| {
239 match item {
240 Ok(event) => serde_json::to_value(event)
241 .map(|value| Envelope::encode_payload(&value))
242 .map_err(|error| {
243 HandlerError::new(
244 ErrorCode::Internal,
245 format!("event serialization failed: {error}"),
246 )
247 }),
248 Err(error) => Err(error.into()),
249 }
250 }))),
251 StreamBody::Raw(stream) => ServiceBody::Stream(stream),
252 };
253 respond(body)
254 }
255}
256
257pub struct State<T>(Arc<T>);
258
259impl<T> State<T> {
260 pub fn new(value: T) -> State<T> {
261 State(Arc::new(value))
262 }
263}
264
265impl<T> Clone for State<T> {
266 fn clone(&self) -> State<T> {
267 State(self.0.clone())
268 }
269}
270
271impl<T> Deref for State<T> {
272 type Target = T;
273
274 fn deref(&self) -> &T {
275 &self.0
276 }
277}
278
279#[derive(Default, Clone)]
280pub(crate) struct StateMap {
281 values: BTreeMap<TypeId, Arc<dyn Any + Send + Sync>>,
282}
283
284impl StateMap {
285 pub(crate) fn insert<T: Send + Sync + 'static>(&mut self, value: T) {
286 self.values.insert(TypeId::of::<T>(), Arc::new(value));
287 }
288
289 pub(crate) fn get<T: Send + Sync + 'static>(&self) -> Option<State<T>> {
290 self.values
291 .get(&TypeId::of::<T>())
292 .cloned()
293 .and_then(|any| any.downcast::<T>().ok())
294 .map(State)
295 }
296
297 pub(crate) fn merged_over(&self, outer: &StateMap) -> StateMap {
298 let mut merged = outer.clone();
299 for (key, value) in &self.values {
300 merged.values.insert(*key, value.clone());
301 }
302 merged
303 }
304}
305
306pub struct States<'a>(pub(crate) &'a StateMap);
307
308impl States<'_> {
309 pub fn state<T: Send + Sync + 'static>(&self) -> Result<State<T>, String> {
310 self.0.get::<T>().ok_or_else(|| {
311 format!(
312 "no registered state provides {}; register it with state(...) on the node or scope",
313 std::any::type_name::<T>()
314 )
315 })
316 }
317}
318
319pub enum ContractSchema {
320 Static(fn() -> Value),
321 Owned(Value),
322}
323
324impl From<fn() -> Value> for ContractSchema {
325 fn from(factory: fn() -> Value) -> ContractSchema {
326 ContractSchema::Static(factory)
327 }
328}
329
330impl From<Value> for ContractSchema {
331 fn from(value: Value) -> ContractSchema {
332 ContractSchema::Owned(value)
333 }
334}
335
336impl ContractSchema {
337 fn value(&self) -> Value {
338 match self {
339 ContractSchema::Static(factory) => factory(),
340 ContractSchema::Owned(value) => value.clone(),
341 }
342 }
343}
344
345pub struct OperationContract {
346 pub input: Option<ContractSchema>,
347 pub output: Option<ContractSchema>,
348 pub event: Option<ContractSchema>,
349 pub error: Option<ContractSchema>,
350}
351
352impl OperationContract {
353 pub fn unknown() -> OperationContract {
354 OperationContract {
355 input: None,
356 output: None,
357 event: None,
358 error: None,
359 }
360 }
361
362 pub(crate) fn to_json(&self, operation: Operation) -> Value {
363 let render = |schema: &Option<ContractSchema>| {
364 schema
365 .as_ref()
366 .map(ContractSchema::value)
367 .unwrap_or_else(|| serde_json::json!({ "unknown": true }))
368 };
369 match operation {
370 Operation::Unary => serde_json::json!({
371 "input_schema": render(&self.input),
372 "output_schema": render(&self.output),
373 }),
374 Operation::Streaming => serde_json::json!({
375 "input_schema": render(&self.input),
376 "event_schema": render(&self.event),
377 "error_schema": render(&self.error),
378 }),
379 }
380 }
381}
382
383type BuildFn = Box<dyn FnOnce(&States<'_>) -> Result<ErasedCall, String> + Send>;
384
385pub struct HandlerService {
386 pub(crate) local_name: String,
387 pub(crate) subject_override: Option<String>,
388 pub(crate) one_line: Option<String>,
389 pub(crate) operation: Operation,
390 pub(crate) metadata: Option<Value>,
391 pub(crate) contract: OperationContract,
392 pub(crate) build: BuildFn,
393}
394
395impl HandlerService {
396 pub fn declare(
397 local_name: &str,
398 subject_override: Option<&str>,
399 one_line: Option<&str>,
400 operation: Operation,
401 contract: OperationContract,
402 build: impl FnOnce(&States<'_>) -> Result<ErasedCall, String> + Send + 'static,
403 ) -> HandlerService {
404 HandlerService {
405 local_name: local_name.into(),
406 subject_override: subject_override.map(Into::into),
407 one_line: one_line.map(Into::into),
408 operation,
409 metadata: None,
410 contract,
411 build: Box::new(build),
412 }
413 }
414
415 pub fn at_subject(mut self, subject: impl Into<String>) -> HandlerService {
416 self.subject_override = Some(subject.into());
417 self
418 }
419
420 pub fn describe(mut self, metadata: Value) -> HandlerService {
421 self.metadata = Some(metadata);
422 self
423 }
424
425 pub(crate) fn effective_subject(&self, scopes: &[String]) -> Result<String, String> {
426 let local = self.subject_override.as_deref().unwrap_or(&self.local_name);
427 let subject = if scopes.is_empty() {
428 local.to_string()
429 } else {
430 format!("{}.{local}", scopes.join("."))
431 };
432 let valid = !subject.is_empty()
433 && subject.len() <= unb_core::MAX_SUBJECT_LEN
434 && subject.split('.').all(|segment| !segment.is_empty());
435 if valid {
436 Ok(subject)
437 } else {
438 Err(format!(
439 "subject {subject:?} needs non-empty dot-separated segments within {} bytes",
440 unb_core::MAX_SUBJECT_LEN
441 ))
442 }
443 }
444}
445
446pub trait Handler: Sized {
447 fn into_service(self) -> HandlerService;
448
449 fn at_subject(self, subject: impl Into<String>) -> HandlerService {
450 self.into_service().at_subject(subject)
451 }
452
453 fn describe(self, metadata: Value) -> HandlerService {
454 self.into_service().describe(metadata)
455 }
456}
457
458impl Handler for HandlerService {
459 fn into_service(self) -> HandlerService {
460 self
461 }
462}
463
464pub fn erase_unary<F, Fut, In, Out, E>(f: F) -> ErasedCall
465where
466 F: Fn(Request<In>) -> Fut + Send + Sync + 'static,
467 Fut: Future<Output = Result<Reply<Out>, E>> + Send + 'static,
468 In: DeserializeOwned + Send + 'static,
469 Out: Serialize + Send + 'static,
470 E: Into<HandlerError> + Send + 'static,
471{
472 let f = Arc::new(f);
473 Arc::new(move |request: http::Request<bytes::Bytes>| {
474 let f = f.clone();
475 Box::pin(async move {
476 let request = Request::<In>::decode(request)?;
477 match f(request).await {
478 Ok(reply) => reply.into_response(),
479 Err(error) => Err(error.into()),
480 }
481 })
482 })
483}
484
485pub fn erase_streaming<F, Fut, In, Event, StreamError, E>(f: F) -> ErasedCall
486where
487 F: Fn(Request<In>) -> Fut + Send + Sync + 'static,
488 Fut: Future<Output = Result<Streaming<Event, StreamError>, E>> + Send + 'static,
489 In: DeserializeOwned + Send + 'static,
490 Event: Serialize + Send + 'static,
491 StreamError: Into<HandlerError> + Send + 'static,
492 E: Into<HandlerError> + Send + 'static,
493{
494 let f = Arc::new(f);
495 Arc::new(move |request: http::Request<bytes::Bytes>| {
496 let f = f.clone();
497 Box::pin(async move {
498 let request = Request::<In>::decode(request)?;
499 match f(request).await {
500 Ok(streaming) => streaming.into_response(),
501 Err(error) => Err(error.into()),
502 }
503 })
504 })
505}
506
507#[cfg(test)]
508mod tests {
509 use serde::Deserialize;
510 use serde_json::json;
511
512 use super::*;
513
514 fn service_request(payload: Value) -> http::Request<bytes::Bytes> {
515 let node = Node::builder("service-test")
516 .insecure_accept_declared_peer_identities()
517 .build()
518 .expect("test node builds");
519 let mut request = http::Request::builder()
520 .method("POST")
521 .uri("/probe")
522 .body(Envelope::encode_payload(&payload))
523 .expect("test request is well formed");
524 request.extensions_mut().insert(Origin::Local);
525 request.extensions_mut().insert(node.snapshot.load_full());
526 request
527 }
528
529 #[derive(Deserialize)]
530 struct Input {
531 n: u32,
532 }
533
534 #[derive(Serialize)]
535 struct Output {
536 doubled: u32,
537 }
538
539 #[tokio::test]
540 async fn a_unary_reply_serializes_through_the_erased_adapter() {
541 let call = erase_unary(|request: Request<Input>| async move {
542 Ok::<_, HandlerError>(Reply::new(Output {
543 doubled: request.payload().n * 2,
544 }))
545 });
546 let response = call(service_request(json!({ "n": 21 })))
547 .await
548 .unwrap_or_else(|error| panic!("call failed: {error}"));
549 let ServiceBody::Unary(payload) = response.into_body() else {
550 panic!("expected a unary response");
551 };
552 let value: Value = serde_json::from_slice(&payload).expect("unary payload is json");
553 assert_eq!(value["doubled"], 42);
554 }
555
556 #[tokio::test]
557 async fn an_undeclared_payload_fails_decode_before_the_handler_runs() {
558 let call = erase_unary(|_request: Request<Input>| async move {
559 panic!("the handler must not run on an invalid payload");
560 #[allow(unreachable_code)]
561 Ok::<_, HandlerError>(Reply::new(Value::Null))
562 });
563 let error = match call(service_request(json!({ "n": "not-a-number" }))).await {
564 Err(error) => error,
565 Ok(_) => panic!("decode must fail"),
566 };
567 assert_eq!(error.code, ErrorCode::InvalidInput);
568 assert!(error.message.contains("probe"));
569 }
570
571 #[tokio::test]
572 async fn a_streaming_output_maps_events_and_errors_into_the_event_stream() {
573 let call = erase_streaming(|_request: Request<Input>| async move {
574 let events = futures_util::stream::iter(vec![
575 Ok(Output { doubled: 2 }),
576 Err(HandlerError::new(ErrorCode::Internal, "stream broke")),
577 ]);
578 Ok::<_, HandlerError>(Streaming::new(events))
579 });
580 let response = call(service_request(json!({ "n": 1 })))
581 .await
582 .unwrap_or_else(|error| panic!("call failed: {error}"));
583 let ServiceBody::Stream(mut stream) = response.into_body() else {
584 panic!("expected a stream response");
585 };
586 let first = stream.next().await.unwrap().unwrap();
587 let first: Value = serde_json::from_slice(&first).unwrap();
588 assert_eq!(first["doubled"], 2);
589 let second = stream.next().await.unwrap().unwrap_err();
590 assert_eq!(second.code, ErrorCode::Internal);
591 assert!(stream.next().await.is_none());
592 }
593
594 #[test]
595 fn nearest_scope_state_wins_over_outer_state() {
596 let mut node_states = StateMap::default();
597 node_states.insert(7u32);
598 node_states.insert("node".to_string());
599 let mut scope_states = StateMap::default();
600 scope_states.insert("scope".to_string());
601 let merged = scope_states.merged_over(&node_states);
602 assert_eq!(*merged.get::<String>().unwrap(), "scope");
603 assert_eq!(*merged.get::<u32>().unwrap(), 7);
604 let missing = match States(&merged).state::<bool>() {
605 Err(missing) => missing,
606 Ok(_) => panic!("bool state must be absent"),
607 };
608 assert!(missing.contains("bool"));
609 }
610}