273bd1b
use std::time::Duration;
use tonic::metadata::{Ascii, MetadataValue};
use tonic::service::Interceptor as TonicInterceptor;
use tonic::transport::{Channel, Endpoint};
use tonic::{Request, Status};
use crate::error::InternalError;
use crate::interceptor::{RPC_KEY_HEADER, new_request_id};
/// 到对端内部 RPC 的共享连接：统一注入 `x-rpc-key` 与请求号。
#[derive(Clone)]
pub struct Connection {
channel: Channel,
interceptor: RpcInterceptor,
}
impl Connection {
pub fn connect(
endpoint: &str,
rpc_key: &str,
timeout: Duration,
) -> Result<Self, InternalError> {
let channel = Endpoint::from_shared(endpoint.to_owned())
.map_err(|error| InternalError::Config(format!("invalid rpc endpoint: {error}")))?
.connect_timeout(timeout)
.timeout(timeout)
.connect_lazy();
Ok(Self {
channel,
interceptor: RpcInterceptor::new(rpc_key)?,
})
}
pub fn channel(&self) -> Channel {
self.channel.clone()
}
pub fn interceptor(&self) -> RpcInterceptor {
self.interceptor.clone()
}
}
#[derive(Clone)]
pub struct RpcInterceptor {
key: Option<MetadataValue<Ascii>>,
request_id: Option<fn() -> Option<String>>,
}
impl RpcInterceptor {
pub fn new(rpc_key: &str) -> Result<Self, InternalError> {
let key =
if rpc_key.trim().is_empty() {
// 空密钥与对端「拒绝一切」语义一致：这里仍然拒绝，避免静默降级。
return Err(InternalError::Config(
"internal rpc key must not be blank".to_string(),
));
} else {
Some(MetadataValue::try_from(rpc_key).map_err(|_| {
InternalError::Config("rpc key is not a valid header".to_string())
})?)
};
Ok(Self {
key,
request_id: None,
})
}
pub fn with_request_id(mut self, provider: fn() -> Option<String>) -> Self {
self.request_id = Some(provider);
self
}
}
impl TonicInterceptor for RpcInterceptor {
fn call(&mut self, mut request: Request<()>) -> Result<Request<()>, Status> {
if let Some(key) = &self.key {
request.metadata_mut().insert(RPC_KEY_HEADER, key.clone());
}
let request_id = self
.request_id
.and_then(|provider| provider())
.unwrap_or_else(new_request_id);
if let Ok(value) = MetadataValue::try_from(request_id) {
request
.metadata_mut()
.insert(crate::interceptor::REQUEST_ID_HEADER, value);
}
Ok(request)
}
}
/// 请求号提供函数别名，方便 `app` 传 `export::middleware::request_id`。
pub type RequestIdProvider = fn() -> Option<String>;