tako_rs_core/grpc/reflection.rs
1//! `grpc.reflection.v1.ServerReflection` scaffolding.
2//!
3//! Provides a [`ReflectionRegistry`](crate::grpc::reflection::ReflectionRegistry) that callers populate with the
4//! `FileDescriptorProto` blobs they want to expose; the registry answers
5//! `list_services`, `file_by_filename`, and `file_containing_symbol` queries.
6//!
7//! ⚠️ **Status:** the encoder/decoder for `ServerReflectionRequest` /
8//! `ServerReflectionResponse` is intentionally minimal — it covers the three
9//! query kinds above and leaves the others (`file_containing_extension`,
10//! `all_extension_numbers_of_type`) to follow-up work. Generated proto code
11//! requires a build script, which would force every consumer to run
12//! `protoc`; the scaffold here lets you ship reflection from a hand-rolled
13//! descriptor or a pre-baked `.pb` blob in the meantime.
14
15use std::sync::Arc;
16
17use parking_lot::RwLock;
18use scc::HashMap as SccHashMap;
19
20/// Reflection registry — populated at startup and consulted by the reflection RPC.
21#[derive(Clone, Default)]
22pub struct ReflectionRegistry {
23 services: Arc<RwLock<Vec<String>>>,
24 files: Arc<SccHashMap<String, Vec<u8>>>,
25 symbols: Arc<SccHashMap<String, String>>,
26}
27
28impl ReflectionRegistry {
29 /// Empty registry.
30 pub fn new() -> Self {
31 Self::default()
32 }
33
34 /// Register a fully-qualified service name (e.g. `helloworld.Greeter`).
35 pub fn add_service(&self, name: impl Into<String>) {
36 self.services.write().push(name.into());
37 }
38
39 /// Register a file descriptor under its source filename.
40 ///
41 /// Re-registering the same filename replaces the prior descriptor (e.g. a
42 /// binary-rebuild that rewrites a `.pb` blob in place).
43 pub fn add_file(&self, filename: impl Into<String>, descriptor: Vec<u8>) {
44 // `upsert_sync` so re-registration of an already-known filename
45 // actually swaps the descriptor — `insert_sync` would silently keep
46 // the stale blob.
47 self.files.upsert_sync(filename.into(), descriptor);
48 }
49
50 /// Map a fully-qualified symbol (`pkg.Service.Method`) to the file that defines it.
51 ///
52 /// Re-mapping an existing symbol replaces the prior filename (e.g. after a
53 /// method moves between files).
54 pub fn map_symbol(&self, symbol: impl Into<String>, filename: impl Into<String>) {
55 // See `add_file`: `upsert_sync` so re-mappings take effect.
56 self.symbols.upsert_sync(symbol.into(), filename.into());
57 }
58
59 /// Snapshot of registered service names.
60 pub fn list_services(&self) -> Vec<String> {
61 self.services.read().clone()
62 }
63
64 /// Look up the descriptor blob for a filename.
65 pub fn file_by_filename(&self, filename: &str) -> Option<Vec<u8>> {
66 self.files.get_sync(filename).map(|e| e.get().clone())
67 }
68
69 /// Resolve a file descriptor by symbol (if registered via `map_symbol`).
70 pub fn file_containing_symbol(&self, symbol: &str) -> Option<Vec<u8>> {
71 let entry = self.symbols.get_sync(symbol)?;
72 let filename = entry.get().clone();
73 drop(entry);
74 self.file_by_filename(&filename)
75 }
76}
77
78/// Marker placed in router state so plugin/middleware code can discover the
79/// active reflection registry.
80pub struct ReflectionState {
81 pub registry: ReflectionRegistry,
82}