Merge pull request #7538 from drmingdrmer/4-echo-ip

feat(meta-service): new RPC to echo client ip

Author: drmingdrmer

src/meta/grpc/src/grpc_action.rs

@@ -16,6 +16,7 @@ use std::convert::TryInto;
 use std::fmt::Debug;
 
 use common_meta_types::protobuf::meta_service_client::MetaServiceClient;
+use common_meta_types::protobuf::ClientInfo;
 use common_meta_types::protobuf::RaftRequest;
 use common_meta_types::protobuf::WatchRequest;
 use common_meta_types::protobuf::WatchResponse;
@@ -35,6 +36,7 @@ use tonic::Request;
 
 use crate::grpc_client::AuthInterceptor;
 use crate::message::ExportReq;
+use crate::message::GetClientInfo;
 use crate::message::GetEndpoints;
 use crate::message::MakeClient;
 
@@ -161,3 +163,7 @@ impl RequestFor for GetEndpoints {
 impl RequestFor for TxnRequest {
     type Reply = TxnReply;
 }
+
+impl RequestFor for GetClientInfo {
+    type Reply = ClientInfo;
+}

src/meta/grpc/src/grpc_client.rs

@@ -41,6 +41,7 @@ use common_grpc::RpcClientTlsConfig;
 use common_meta_api::KVApi;
 use common_meta_types::anyerror::AnyError;
 use common_meta_types::protobuf::meta_service_client::MetaServiceClient;
+use common_meta_types::protobuf::ClientInfo;
 use common_meta_types::protobuf::Empty;
 use common_meta_types::protobuf::ExportedChunk;
 use common_meta_types::protobuf::HandshakeRequest;
@@ -213,6 +214,10 @@ impl ClientHandle {
         Ok(r)
     }
 
+    pub async fn get_client_info(&self) -> std::result::Result<ClientInfo, MetaError> {
+        self.request(message::GetClientInfo {}).await
+    }
+
     pub async fn make_client(
         &self,
     ) -> std::result::Result<
@@ -379,6 +384,10 @@ impl MetaGrpcClient {
                     let resp = self.get_endpoints().await;
                     Ok(message::Response::GetEndpoints(resp))
                 }
+                message::Request::GetClientInfo(_) => {
+                    let resp = self.get_client_info().await;
+                    resp.map(message::Response::GetClientInfo)
+                }
             };
 
             debug!(
@@ -752,6 +761,16 @@ impl MetaGrpcClient {
         Ok(res.into_inner())
     }
 
+    /// Export all data in json from metasrv.
+    #[tracing::instrument(level = "debug", skip_all)]
+    pub(crate) async fn get_client_info(&self) -> std::result::Result<ClientInfo, MetaError> {
+        debug!("MetaGrpcClient::get_client_info");
+
+        let mut client = self.make_client().await?;
+        let res = client.get_client_info(Empty {}).await?;
+        Ok(res.into_inner())
+    }
+
     #[tracing::instrument(level = "debug", skip(self, v))]
     pub(crate) async fn do_write<T, R>(&self, v: T) -> std::result::Result<R, MetaError>
     where

src/meta/grpc/src/message.rs

@@ -14,6 +14,7 @@
 
 use common_base::base::tokio::sync::oneshot::Sender;
 use common_meta_types::protobuf::meta_service_client::MetaServiceClient;
+use common_meta_types::protobuf::ClientInfo;
 use common_meta_types::protobuf::ExportedChunk;
 use common_meta_types::protobuf::WatchRequest;
 use common_meta_types::protobuf::WatchResponse;
@@ -72,6 +73,9 @@ pub enum Request {
 
     /// Get endpoints, for test
     GetEndpoints(GetEndpoints),
+
+    /// Get info about the client
+    GetClientInfo(GetClientInfo),
 }
 
 /// Meta-client worker-to-handle response body
@@ -86,6 +90,7 @@ pub enum Response {
     Export(tonic::codec::Streaming<ExportedChunk>),
     MakeClient(MetaServiceClient<InterceptedService<Channel, AuthInterceptor>>),
     GetEndpoints(Vec<String>),
+    GetClientInfo(ClientInfo),
 }
 
 /// Export all data stored in metasrv
@@ -101,3 +106,7 @@ pub struct MakeClient {}
 /// Get all meta server endpoints
 #[derive(serde::Serialize, serde::Deserialize, Clone, Debug)]
 pub struct GetEndpoints {}
+
+/// Get info about client
+#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)]
+pub struct GetClientInfo {}

src/meta/grpc/tests/it/grpc_server.rs

@@ -19,6 +19,8 @@ use std::time::Duration;
 use common_base::base::tokio;
 use common_meta_types::protobuf::meta_service_server::MetaService;
 use common_meta_types::protobuf::meta_service_server::MetaServiceServer;
+use common_meta_types::protobuf::ClientInfo;
+use common_meta_types::protobuf::Empty;
 use common_meta_types::protobuf::ExportedChunk;
 use common_meta_types::protobuf::HandshakeResponse;
 use common_meta_types::protobuf::MemberListReply;
@@ -102,6 +104,13 @@ impl MetaService for GrpcServiceForTestImpl {
     ) -> Result<Response<MemberListReply>, Status> {
         todo!()
     }
+
+    async fn get_client_info(
+        &self,
+        _request: Request<Empty>,
+    ) -> Result<Response<ClientInfo>, Status> {
+        todo!()
+    }
 }
 
 pub fn start_grpc_server() -> String {

src/meta/service/src/api/grpc/grpc_service.rs

@@ -24,6 +24,8 @@ use common_grpc::GrpcToken;
 use common_meta_grpc::MetaGrpcReadReq;
 use common_meta_grpc::MetaGrpcWriteReq;
 use common_meta_types::protobuf::meta_service_server::MetaService;
+use common_meta_types::protobuf::ClientInfo;
+use common_meta_types::protobuf::Empty;
 use common_meta_types::protobuf::ExportedChunk;
 use common_meta_types::protobuf::HandshakeRequest;
 use common_meta_types::protobuf::HandshakeResponse;
@@ -264,6 +266,20 @@ impl MetaService for MetaServiceImpl {
 
         Ok(Response::new(resp))
     }
+
+    async fn get_client_info(
+        &self,
+        request: Request<Empty>,
+    ) -> Result<Response<ClientInfo>, Status> {
+        let r = request.remote_addr();
+        if let Some(addr) = r {
+            let resp = ClientInfo {
+                client_addr: addr.to_string(),
+            };
+            return Ok(Response::new(resp));
+        }
+        Err(Status::unavailable("can not get client ip address"))
+    }
 }
 
 pub struct ExportStream {

src/meta/service/tests/it/grpc/metasrv_grpc_get_client_info.rs

@@ -0,0 +1,51 @@
+// Copyright 2022 Datafuse Labs.
+//
+// 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::time::Duration;
+
+use common_base::base::tokio;
+use common_meta_grpc::MetaGrpcClient;
+use pretty_assertions::assert_eq;
+use regex::Regex;
+
+use crate::init_meta_ut;
+
+#[async_entry::test(worker_threads = 3, init = "init_meta_ut!()", tracing_span = "debug")]
+async fn test_get_client_info() -> anyhow::Result<()> {
+    // - Start a metasrv server.
+    // - Get client ip
+
+    let (_tc, addr) = crate::tests::start_metasrv().await?;
+
+    let client = MetaGrpcClient::try_create(
+        vec![addr],
+        "root",
+        "xxx",
+        None,
+        Some(Duration::from_secs(10)),
+        None,
+    )?;
+
+    let resp = client.get_client_info().await?;
+
+    let client_addr = resp.client_addr;
+
+    let masked_addr = Regex::new(r"\d+")
+        .unwrap()
+        .replace_all(&client_addr, "1")
+        .to_string();
+
+    assert_eq!("1.1.1.1:1", masked_addr);
+    Ok(())
+}

src/meta/service/tests/it/grpc/mod.rs

@@ -14,6 +14,7 @@
 
 pub mod metasrv_grpc_api;
 mod metasrv_grpc_export;
+pub mod metasrv_grpc_get_client_info;
 pub mod metasrv_grpc_handshake;
 pub mod metasrv_grpc_kv_api;
 pub mod metasrv_grpc_kv_api_restart_cluster;

src/meta/types/proto/meta.proto

@@ -143,6 +143,11 @@ message TxnReply {
   string error = 3;
 }
 
+message ClientInfo {
+  // The address of the connected in form of "<ip>:<port>"
+  string client_addr = 10;
+}
+
 service RaftService {
 
   rpc Write(RaftRequest) returns (RaftReply) {}
@@ -181,4 +186,7 @@ service MetaService {
 
   // Get MetaSrv member list endpoints
   rpc MemberList(MemberListRequest) returns (MemberListReply);
+
+  // Respond with the information about the client.
+  rpc GetClientInfo(Empty) returns (ClientInfo);
 }