Refactor RPC Authentication System for Improved Maintainability (#1391)

2026-09-04 11:15:39 +00:00 · 2026-01-05 19:51:51 +08:00
parent 0b6f3302ce
commit 5660208e89
15 changed files with 697 additions and 299 deletions
@@ -0,0 +1,88 @@
+// Copyright 2024 RustFS Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::error::Error;
+
+use http::Method;
+use rustfs_common::GLOBAL_CONN_MAP;
+use rustfs_protos::{create_new_channel, proto_gen::node_service::node_service_client::NodeServiceClient};
+use tonic::{service::interceptor::InterceptedService, transport::Channel};
+use tracing::debug;
+
+use crate::rpc::{TONIC_RPC_PREFIX, gen_signature_headers};
+
+/// 3. Subsequent calls will attempt fresh connections
+/// 4. If node is still down, connection will fail fast (3s timeout)
+pub async fn node_service_time_out_client(
+    addr: &String,
+    interceptor: TonicInterceptor,
+) -> Result<NodeServiceClient<InterceptedService<Channel, TonicInterceptor>>, Box<dyn Error>> {
+    // Try to get cached channel
+    let cached_channel = { GLOBAL_CONN_MAP.read().await.get(addr).cloned() };
+
+    let channel = match cached_channel {
+        Some(channel) => {
+            debug!("Using cached gRPC channel for: {}", addr);
+            channel
+        }
+        None => {
+            // No cached connection, create new one
+            create_new_channel(addr).await?
+        }
+    };
+
+    Ok(NodeServiceClient::with_interceptor(channel, interceptor))
+}
+
+pub async fn node_service_time_out_client_no_auth(
+    addr: &String,
+) -> Result<NodeServiceClient<InterceptedService<Channel, TonicInterceptor>>, Box<dyn Error>> {
+    node_service_time_out_client(addr, TonicInterceptor::NoOp(NoOpInterceptor)).await
+}
+
+pub struct TonicSignatureInterceptor;
+
+impl tonic::service::Interceptor for TonicSignatureInterceptor {
+    fn call(&mut self, mut req: tonic::Request<()>) -> Result<tonic::Request<()>, tonic::Status> {
+        let headers = gen_signature_headers(TONIC_RPC_PREFIX, &Method::GET);
+        req.metadata_mut().as_mut().extend(headers);
+        Ok(req)
+    }
+}
+
+pub fn gen_tonic_signature_interceptor() -> TonicSignatureInterceptor {
+    TonicSignatureInterceptor
+}
+
+pub struct NoOpInterceptor;
+
+impl tonic::service::Interceptor for NoOpInterceptor {
+    fn call(&mut self, req: tonic::Request<()>) -> Result<tonic::Request<()>, tonic::Status> {
+        Ok(req)
+    }
+}
+
+pub enum TonicInterceptor {
+    Signature(TonicSignatureInterceptor),
+    NoOp(NoOpInterceptor),
+}
+
+impl tonic::service::Interceptor for TonicInterceptor {
+    fn call(&mut self, req: tonic::Request<()>) -> Result<tonic::Request<()>, tonic::Status> {
+        match self {
+            TonicInterceptor::Signature(interceptor) => interceptor.call(req),
+            TonicInterceptor::NoOp(interceptor) => interceptor.call(req),
+        }
+    }
+}