tonic_reflection/server/
v1.rs1use std::{fmt, sync::Arc};
26
27use tokio::sync::mpsc;
28use tokio_stream::{Stream, StreamExt};
29use tonic::{Request, Response, Status, Streaming};
30
31use super::ReflectionServiceState;
32use crate::pb::v1::server_reflection_request::MessageRequest;
33use crate::pb::v1::server_reflection_response::MessageResponse;
34pub use crate::pb::v1::server_reflection_server::{ServerReflection, ServerReflectionServer};
35use crate::pb::v1::{
36 ExtensionNumberResponse, FileDescriptorResponse, ListServiceResponse, ServerReflectionRequest,
37 ServerReflectionResponse, ServiceResponse,
38};
39
40#[derive(Debug)]
42pub struct ReflectionService {
43 state: Arc<ReflectionServiceState>,
44}
45
46#[tonic::async_trait]
47impl ServerReflection for ReflectionService {
48 type ServerReflectionInfoStream = ServerReflectionInfoStream;
49
50 async fn server_reflection_info(
51 &self,
52 req: Request<Streaming<ServerReflectionRequest>>,
53 ) -> Result<Response<Self::ServerReflectionInfoStream>, Status> {
54 let mut req_rx = req.into_inner();
55 let (resp_tx, resp_rx) = mpsc::channel::<Result<ServerReflectionResponse, Status>>(1);
56
57 let state = self.state.clone();
58
59 tokio::spawn(async move {
60 while let Some(req) = req_rx.next().await {
61 let Ok(req) = req else {
62 return;
63 };
64
65 let resp_msg = match req.message_request.clone() {
66 None => Err(Status::invalid_argument("invalid MessageRequest")),
67 Some(msg) => match msg {
68 MessageRequest::FileByFilename(s) => state.file_by_filename(&s).map(|fd| {
69 MessageResponse::FileDescriptorResponse(FileDescriptorResponse {
70 file_descriptor_proto: vec![fd],
71 })
72 }),
73 MessageRequest::FileContainingSymbol(s) => {
74 state.symbol_by_name(&s).map(|fd| {
75 MessageResponse::FileDescriptorResponse(FileDescriptorResponse {
76 file_descriptor_proto: vec![fd],
77 })
78 })
79 }
80 MessageRequest::FileContainingExtension(_) => {
81 Err(Status::not_found("extensions are not supported"))
82 }
83 MessageRequest::AllExtensionNumbersOfType(_) => {
84 Ok(MessageResponse::AllExtensionNumbersResponse(
87 ExtensionNumberResponse::default(),
88 ))
89 }
90 MessageRequest::ListServices(_) => {
91 Ok(MessageResponse::ListServicesResponse(ListServiceResponse {
92 service: state
93 .list_services()
94 .iter()
95 .map(|s| ServiceResponse { name: s.clone() })
96 .collect(),
97 }))
98 }
99 },
100 };
101
102 match resp_msg {
103 Ok(resp_msg) => {
104 let resp = ServerReflectionResponse {
105 valid_host: req.host.clone(),
106 original_request: Some(req.clone()),
107 message_response: Some(resp_msg),
108 };
109 if resp_tx.send(Ok(resp)).await.is_err() {
110 return;
111 }
112 }
113 Err(status) => {
114 let _ = resp_tx.send(Err(status)).await;
115 return;
116 }
117 }
118 }
119 });
120
121 Ok(Response::new(ServerReflectionInfoStream::new(resp_rx)))
122 }
123}
124
125impl From<ReflectionServiceState> for ReflectionService {
126 fn from(state: ReflectionServiceState) -> Self {
127 Self {
128 state: Arc::new(state),
129 }
130 }
131}
132
133pub struct ServerReflectionInfoStream {
135 inner: tokio_stream::wrappers::ReceiverStream<Result<ServerReflectionResponse, Status>>,
136}
137
138impl ServerReflectionInfoStream {
139 fn new(resp_rx: mpsc::Receiver<Result<ServerReflectionResponse, Status>>) -> Self {
140 let inner = tokio_stream::wrappers::ReceiverStream::new(resp_rx);
141 Self { inner }
142 }
143}
144
145impl Stream for ServerReflectionInfoStream {
146 type Item = Result<ServerReflectionResponse, Status>;
147
148 fn poll_next(
149 mut self: std::pin::Pin<&mut Self>,
150 cx: &mut std::task::Context<'_>,
151 ) -> std::task::Poll<Option<Self::Item>> {
152 std::pin::Pin::new(&mut self.inner).poll_next(cx)
153 }
154
155 fn size_hint(&self) -> (usize, Option<usize>) {
156 self.inner.size_hint()
157 }
158}
159
160impl fmt::Debug for ServerReflectionInfoStream {
161 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
162 f.debug_tuple("ServerReflectionInfoStream").finish()
163 }
164}