main
use std::collections::{BTreeMap, HashSet};
use std::pin::Pin;
use std::sync::Arc;
use prost::Message;
use proto::grpc::reflection::v1::server_reflection_request::MessageRequest;
use proto::grpc::reflection::v1::server_reflection_response::MessageResponse;
use proto::grpc::reflection::v1::server_reflection_server::{
ServerReflection as ReflectionTrait, ServerReflectionServer,
};
use proto::grpc::reflection::v1::{
ErrorResponse, FileDescriptorResponse, ListServiceResponse, ServerReflectionRequest,
ServerReflectionResponse, ServiceResponse,
};
use tonic::{Request, Response, Status};
/// 基于 `lib/proto` 的 FILE_DESCRIPTOR_SET 实现标准 gRPC 反射（v1）。
///
/// 不引入 tonic-reflection：把 descriptor 解码成 FileDescriptorProto 后自行按
/// list_services / file_by_filename / file_containing_symbol 响应。
#[derive(Clone)]
pub struct ReflectionService {
files: Arc<Vec<prost_types::FileDescriptorProto>>,
by_name: Arc<BTreeMap<String, usize>>,
services: Arc<Vec<String>>,
}
impl ReflectionService {
pub fn new(service_names: &[&'static str]) -> Result<Self, crate::error::InternalError> {
let descriptor = prost_types::FileDescriptorSet::decode(proto::FILE_DESCRIPTOR_SET)
.map_err(|error| {
crate::error::InternalError::Config(format!(
"decoding proto descriptor set: {error}"
))
})?;
let files = descriptor.file;
let by_name = files
.iter()
.enumerate()
.filter_map(|(index, file)| file.name.clone().map(|name| (name, index)))
.collect();
let mut services = Vec::new();
for name in service_names {
services.push((*name).to_string());
}
for name in [
crate::health::HEALTH_SERVICE,
crate::health::REFLECTION_SERVICE,
] {
if !services.iter().any(|service| service == name) {
services.push(name.to_string());
}
}
Ok(Self {
files: Arc::new(files),
by_name: Arc::new(by_name),
services: Arc::new(services),
})
}
pub fn server(self) -> ServerReflectionServer<Self> {
ServerReflectionServer::new(self)
}
fn response(
&self,
request: ServerReflectionRequest,
message_response: MessageResponse,
) -> ServerReflectionResponse {
ServerReflectionResponse {
valid_host: request.host.clone(),
original_request: Some(request),
message_response: Some(message_response),
}
}
fn error(
&self,
request: ServerReflectionRequest,
code: i32,
message: &str,
) -> ServerReflectionResponse {
self.response(
request,
MessageResponse::ErrorResponse(ErrorResponse {
error_code: code,
error_message: message.to_string(),
}),
)
}
fn file_response(
&self,
request: ServerReflectionRequest,
roots: Vec<usize>,
) -> ServerReflectionResponse {
let mut seen = HashSet::new();
let mut queue = roots;
let mut encoded = Vec::new();
while let Some(index) = queue.pop() {
if !seen.insert(index) {
continue;
}
let Some(file) = self.files.get(index) else {
continue;
};
encoded.push(file.encode_to_vec());
for dependency in &file.dependency {
if let Some(next) = self.by_name.get(dependency) {
queue.push(*next);
}
}
}
self.response(
request,
MessageResponse::FileDescriptorResponse(FileDescriptorResponse {
file_descriptor_proto: encoded,
}),
)
}
fn handle(&self, request: ServerReflectionRequest) -> ServerReflectionResponse {
match request.message_request.clone() {
Some(MessageRequest::ListServices(_)) => {
let services = self
.services
.iter()
.map(|name| ServiceResponse { name: name.clone() })
.collect();
self.response(
request,
MessageResponse::ListServicesResponse(ListServiceResponse {
service: services,
}),
)
}
Some(MessageRequest::FileByFilename(name)) => {
let Some(index) = self.by_name.get(&name).copied() else {
return self.error(request, 5, "file not found");
};
self.file_response(request, vec![index])
}
Some(MessageRequest::FileContainingSymbol(symbol)) => {
match self.find_file_by_symbol(&symbol) {
Some(index) => self.file_response(request, vec![index]),
None => self.error(request, 5, "symbol not found"),
}
}
Some(MessageRequest::AllExtensionNumbersOfType(_))
| Some(MessageRequest::FileContainingExtension(_)) => {
self.error(request, 12, "extension reflection is not implemented")
}
None => self.error(request, 3, "invalid reflection request"),
}
}
fn find_file_by_symbol(&self, symbol: &str) -> Option<usize> {
let symbol = symbol.trim();
for (index, file) in self.files.iter().enumerate() {
if file.package.as_deref() == Some(symbol) {
return Some(index);
}
let package = file.package.as_deref().unwrap_or_default();
for service in &file.service {
let Some(name) = service.name.as_deref() else {
continue;
};
let full = if package.is_empty() {
name.to_string()
} else {
format!("{package}.{name}")
};
if full == symbol {
return Some(index);
}
if let Some(method) = symbol
.strip_prefix(&full)
.and_then(|rest| rest.strip_prefix('.'))
{
if service
.method
.iter()
.any(|item| item.name.as_deref() == Some(method))
{
return Some(index);
}
}
}
for message in &file.message_type {
if let Some(name) = message.name.as_deref() {
let full = if package.is_empty() {
name.to_string()
} else {
format!("{package}.{name}")
};
if full == symbol {
return Some(index);
}
}
}
}
None
}
}
#[tonic::async_trait]
impl ReflectionTrait for ReflectionService {
type ServerReflectionInfoStream = Pin<
Box<
dyn tokio_stream::Stream<Item = Result<ServerReflectionResponse, Status>>
+ Send
+ 'static,
>,
>;
async fn server_reflection_info(
&self,
request: Request<tonic::Streaming<ServerReflectionRequest>>,
) -> Result<Response<Self::ServerReflectionInfoStream>, Status> {
let mut inbound = request.into_inner();
let this = self.clone();
let (tx, rx) = tokio::sync::mpsc::channel(8);
tokio::spawn(async move {
loop {
match tokio_stream::StreamExt::next(&mut inbound).await {
None => break,
Some(Ok(request)) => {
if tx.send(Ok(this.handle(request))).await.is_err() {
break;
}
}
Some(Err(_)) => break,
}
}
});
Ok(Response::new(Box::pin(
tokio_stream::wrappers::ReceiverStream::new(rx),
)))
}
}