Skip to main content

tonic_reflection/server/
v1alpha.rs

1/*
2 *
3 * Copyright 2025 gRPC authors.
4 *
5 * Permission is hereby granted, free of charge, to any person obtaining a copy
6 * of this software and associated documentation files (the "Software"), to
7 * deal in the Software without restriction, including without limitation the
8 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
9 * sell copies of the Software, and to permit persons to whom the Software is
10 * furnished to do so, subject to the following conditions:
11 *
12 * The above copyright notice and this permission notice shall be included in
13 * all copies or substantial portions of the Software.
14 *
15 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
20 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
21 * IN THE SOFTWARE.
22 *
23 */
24
25use 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::v1alpha::server_reflection_request::MessageRequest;
33use crate::pb::v1alpha::server_reflection_response::MessageResponse;
34pub use crate::pb::v1alpha::server_reflection_server::{ServerReflection, ServerReflectionServer};
35use crate::pb::v1alpha::{
36    ExtensionNumberResponse, FileDescriptorResponse, ListServiceResponse, ServerReflectionRequest,
37    ServerReflectionResponse, ServiceResponse,
38};
39
40/// An implementation for `ServerReflection`.
41#[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                            // NOTE: Workaround. Some grpc clients (e.g. grpcurl) expect this method not to fail.
85                            // https://github.com/grpc/grpc-rust/issues/1077
86                            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
133/// A response stream.
134pub 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}