Skip to main content

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}