tonic-reflection 0.14.9

/*
 *
 * Copyright 2025 gRPC authors.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 * IN THE SOFTWARE.
 *
 */

use std::collections::HashMap;
use std::fmt::{Display, Formatter};
use std::sync::Arc;

use prost::encoding::{WireType, decode_key, decode_varint};
use prost::{DecodeError, Message};
use prost_types::{
    DescriptorProto, EnumDescriptorProto, FieldDescriptorProto, FileDescriptorProto,
    FileDescriptorSet,
};
use tonic::Status;

/// v1 interface for the gRPC Reflection Service server.
pub mod v1;
/// v1alpha interface for the gRPC Reflection Service server.
pub mod v1alpha;

/// A builder used to construct a gRPC Reflection Service.
#[derive(Debug)]
pub struct Builder<'b> {
    file_descriptor_sets: Vec<FileDescriptorSet>,
    encoded_file_descriptor_sets: Vec<&'b [u8]>,
    include_reflection_service: bool,

    service_names: Vec<String>,
    use_all_service_names: bool,
}

impl<'b> Builder<'b> {
    /// Create a new builder that can configure a gRPC Reflection Service.
    pub fn configure() -> Self {
        Builder {
            file_descriptor_sets: Vec::new(),
            encoded_file_descriptor_sets: Vec::new(),
            include_reflection_service: true,

            service_names: Vec::new(),
            use_all_service_names: true,
        }
    }

    /// Registers an instance of `prost_types::FileDescriptorSet` with the gRPC Reflection
    /// Service builder.
    pub fn register_file_descriptor_set(mut self, file_descriptor_set: FileDescriptorSet) -> Self {
        self.file_descriptor_sets.push(file_descriptor_set);
        self
    }

    /// Registers a byte slice containing an encoded `prost_types::FileDescriptorSet` with
    /// the gRPC Reflection Service builder.
    pub fn register_encoded_file_descriptor_set(
        mut self,
        encoded_file_descriptor_set: &'b [u8],
    ) -> Self {
        self.encoded_file_descriptor_sets
            .push(encoded_file_descriptor_set);
        self
    }

    /// Serve the gRPC Reflection Service descriptor via the Reflection Service. This is enabled
    /// by default - set `include` to false to disable.
    pub fn include_reflection_service(mut self, include: bool) -> Self {
        self.include_reflection_service = include;
        self
    }

    /// Advertise a fully-qualified gRPC service name.
    ///
    /// If not called, then all services present in the registered file descriptor sets
    /// will be advertised.
    pub fn with_service_name(mut self, name: impl Into<String>) -> Self {
        self.use_all_service_names = false;
        self.service_names.push(name.into());
        self
    }

    /// Build a v1 gRPC Reflection Service to be served via Tonic.
    pub fn build_v1(
        mut self,
    ) -> Result<v1::ServerReflectionServer<impl v1::ServerReflection>, Error> {
        if self.include_reflection_service {
            self = self.register_encoded_file_descriptor_set(crate::pb::v1::FILE_DESCRIPTOR_SET);
        }

        Ok(v1::ServerReflectionServer::new(
            v1::ReflectionService::from(ReflectionServiceState::new(
                self.service_names,
                self.encoded_file_descriptor_sets,
                self.file_descriptor_sets,
                self.use_all_service_names,
            )?),
        ))
    }

    /// Build a v1alpha gRPC Reflection Service to be served via Tonic.
    pub fn build_v1alpha(
        mut self,
    ) -> Result<v1alpha::ServerReflectionServer<impl v1alpha::ServerReflection>, Error> {
        if self.include_reflection_service {
            self =
                self.register_encoded_file_descriptor_set(crate::pb::v1alpha::FILE_DESCRIPTOR_SET);
        }

        Ok(v1alpha::ServerReflectionServer::new(
            v1alpha::ReflectionService::from(ReflectionServiceState::new(
                self.service_names,
                self.encoded_file_descriptor_sets,
                self.file_descriptor_sets,
                self.use_all_service_names,
            )?),
        ))
    }
}

#[derive(Debug)]
struct ReflectionServiceState {
    service_names: Vec<String>,
    // The original, encoded bytes of each `FileDescriptorProto` are stored (and
    // served) verbatim rather than a decoded `prost_types::FileDescriptorProto`.
    // prost silently drops unknown fields on decode, so a decode/re-encode
    // round-trip strips custom options / extensions (e.g. `google.api.http`);
    // keeping the raw bytes preserves them for reflection consumers.
    files: HashMap<String, Arc<[u8]>>,
    symbols: HashMap<String, Arc<[u8]>>,
}

impl ReflectionServiceState {
    fn new(
        service_names: Vec<String>,
        encoded_file_descriptor_sets: Vec<&[u8]>,
        file_descriptor_sets: Vec<FileDescriptorSet>,
        use_all_service_names: bool,
    ) -> Result<Self, Error> {
        let mut state = ReflectionServiceState {
            service_names,
            files: HashMap::new(),
            symbols: HashMap::new(),
        };

        // The builder accepts descriptor sets in two forms: raw bytes (via
        // `register_encoded_file_descriptor_set`) and already-decoded
        // `prost_types::FileDescriptorSet`s (via `register_file_descriptor_set`).
        //
        // The encoded sets are registered first, on purpose: `register_file`
        // keeps the first entry seen for a given file name (see below), and the
        // encoded form is the one that faithfully preserves custom options. So
        // if the same file arrives in both forms, the lossless bytes win.
        //
        // Each encoded set is split into its per-file byte ranges without
        // decoding the inner message, so the original bytes (including custom
        // options such as `google.api.http`) are retained and served unchanged.
        for encoded in encoded_file_descriptor_sets {
            for raw in split_file_descriptor_set(encoded)? {
                let fd = FileDescriptorProto::decode(raw.as_slice())?;
                state.register_file(&fd, Arc::from(raw), use_all_service_names)?;
            }
        }

        // Already-decoded sets carry no original bytes, so each file has to be
        // re-encoded. This is lossy for custom options, but unavoidable: those
        // options were already dropped when the caller decoded the set into
        // `prost_types` (which has no field to hold them). Callers that need
        // custom options preserved should use `register_encoded_file_descriptor_set`
        // instead. Registered second, as a fallback for files not already
        // supplied in encoded form above.
        for fds in file_descriptor_sets {
            for fd in fds.file {
                let raw = fd.encode_to_vec();
                state.register_file(&fd, Arc::from(raw), use_all_service_names)?;
            }
        }

        Ok(state)
    }

    fn register_file(
        &mut self,
        fd: &FileDescriptorProto,
        raw: Arc<[u8]>,
        use_all_service_names: bool,
    ) -> Result<(), Error> {
        let name = match fd.name.clone() {
            None => {
                return Err(Error::InvalidFileDescriptorSet("missing name".to_string()));
            }
            Some(n) => n,
        };

        // First registration for a file name wins; later duplicates are ignored.
        // `new` registers the option-preserving encoded sets before the
        // re-encoded ones, so a file present in both keeps its original bytes.
        if self.files.contains_key(&name) {
            return Ok(());
        }

        self.files.insert(name, raw.clone());
        self.process_file(fd, &raw, use_all_service_names)
    }

    fn process_file(
        &mut self,
        fd: &FileDescriptorProto,
        raw: &Arc<[u8]>,
        use_all_service_names: bool,
    ) -> Result<(), Error> {
        let prefix = &fd.package.clone().unwrap_or_default();

        for msg in &fd.message_type {
            self.process_message(raw, prefix, msg)?;
        }

        for en in &fd.enum_type {
            self.process_enum(raw, prefix, en)?;
        }

        for service in &fd.service {
            let service_name = extract_name(prefix, "service", service.name.as_ref())?;
            if use_all_service_names {
                self.service_names.push(service_name.clone());
            }
            self.symbols.insert(service_name.clone(), raw.clone());

            for method in &service.method {
                let method_name = extract_name(&service_name, "method", method.name.as_ref())?;
                self.symbols.insert(method_name, raw.clone());
            }
        }

        Ok(())
    }

    fn process_message(
        &mut self,
        raw: &Arc<[u8]>,
        prefix: &str,
        msg: &DescriptorProto,
    ) -> Result<(), Error> {
        let message_name = extract_name(prefix, "message", msg.name.as_ref())?;
        self.symbols.insert(message_name.clone(), raw.clone());

        for nested in &msg.nested_type {
            self.process_message(raw, &message_name, nested)?;
        }

        for en in &msg.enum_type {
            self.process_enum(raw, &message_name, en)?;
        }

        for field in &msg.field {
            self.process_field(raw, &message_name, field)?;
        }

        for oneof in &msg.oneof_decl {
            let oneof_name = extract_name(&message_name, "oneof", oneof.name.as_ref())?;
            self.symbols.insert(oneof_name, raw.clone());
        }

        Ok(())
    }

    fn process_enum(
        &mut self,
        raw: &Arc<[u8]>,
        prefix: &str,
        en: &EnumDescriptorProto,
    ) -> Result<(), Error> {
        let enum_name = extract_name(prefix, "enum", en.name.as_ref())?;
        self.symbols.insert(enum_name.clone(), raw.clone());

        for value in &en.value {
            let value_name = extract_name(&enum_name, "enum value", value.name.as_ref())?;
            self.symbols.insert(value_name, raw.clone());
        }

        Ok(())
    }

    fn process_field(
        &mut self,
        raw: &Arc<[u8]>,
        prefix: &str,
        field: &FieldDescriptorProto,
    ) -> Result<(), Error> {
        let field_name = extract_name(prefix, "field", field.name.as_ref())?;
        self.symbols.insert(field_name, raw.clone());
        Ok(())
    }

    fn list_services(&self) -> &[String] {
        &self.service_names
    }

    fn symbol_by_name(&self, symbol: &str) -> Result<Vec<u8>, Status> {
        match self.symbols.get(symbol) {
            None => Err(Status::not_found(format!("symbol '{symbol}' not found"))),
            Some(raw) => Ok(raw.to_vec()),
        }
    }

    fn file_by_filename(&self, filename: &str) -> Result<Vec<u8>, Status> {
        match self.files.get(filename) {
            None => Err(Status::not_found(format!("file '{filename}' not found"))),
            Some(raw) => Ok(raw.to_vec()),
        }
    }
}

/// Splits an encoded `FileDescriptorSet` into the raw, unmodified bytes of each
/// contained `FileDescriptorProto` (wire field 1), without decoding the inner
/// messages.
///
/// Serving these original bytes preserves custom options / extensions (e.g.
/// `google.api.http`) that prost silently drops when a `FileDescriptorProto` is
/// decoded and re-encoded.
fn split_file_descriptor_set(mut buf: &[u8]) -> Result<Vec<Vec<u8>>, Error> {
    // Field number of `repeated FileDescriptorProto file = 1;`.
    const FILE_FIELD_TAG: u32 = 1;

    let truncated = || {
        Error::InvalidFileDescriptorSet("unexpected end of FileDescriptorSet buffer".to_string())
    };

    let mut files = Vec::new();
    while !buf.is_empty() {
        let (tag, wire_type) = decode_key(&mut buf)?;
        match wire_type {
            WireType::LengthDelimited => {
                let len = decode_varint(&mut buf)? as usize;
                if len > buf.len() {
                    return Err(truncated());
                }
                let (field, rest) = buf.split_at(len);
                if tag == FILE_FIELD_TAG {
                    files.push(field.to_vec());
                }
                buf = rest;
            }
            WireType::Varint => {
                decode_varint(&mut buf)?;
            }
            WireType::ThirtyTwoBit => {
                if buf.len() < 4 {
                    return Err(truncated());
                }
                buf = &buf[4..];
            }
            WireType::SixtyFourBit => {
                if buf.len() < 8 {
                    return Err(truncated());
                }
                buf = &buf[8..];
            }
            WireType::StartGroup | WireType::EndGroup => {
                return Err(Error::InvalidFileDescriptorSet(
                    "unexpected group wire type in FileDescriptorSet".to_string(),
                ));
            }
        }
    }

    Ok(files)
}

fn extract_name(
    prefix: &str,
    name_type: &str,
    maybe_name: Option<&String>,
) -> Result<String, Error> {
    match maybe_name {
        None => Err(Error::InvalidFileDescriptorSet(format!(
            "missing {name_type} name"
        ))),
        Some(name) => {
            if prefix.is_empty() {
                Ok(name.to_string())
            } else {
                Ok(format!("{prefix}.{name}"))
            }
        }
    }
}

/// Represents an error in the construction of a gRPC Reflection Service.
#[derive(Debug)]
pub enum Error {
    /// An error was encountered decoding a `prost_types::FileDescriptorSet` from a buffer.
    DecodeError(prost::DecodeError),
    /// An invalid `prost_types::FileDescriptorProto` was encountered.
    InvalidFileDescriptorSet(String),
}

impl From<DecodeError> for Error {
    fn from(e: DecodeError) -> Self {
        Error::DecodeError(e)
    }
}

impl std::error::Error for Error {}

impl Display for Error {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Error::DecodeError(_) => f.write_str("error decoding FileDescriptorSet from buffer"),
            Error::InvalidFileDescriptorSet(s) => {
                write!(f, "invalid FileDescriptorSet - {s}")
            }
        }
    }
}