add callee ip in RpcContext (#65)

* feat: rename contextWithRemoteIp to contextWithRemoteCallerIp * feat: add callee ip in context before invoke rpc request * feat: add unit test * feat:fix * import issue * feat:add caller/callee port in RpcContext
trpc-group · Sep 23, 2024 · 9018c10 · 9018c10
1 parent acdbb07
commit 9018c10
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 8 deletions.
diff --git a/trpc-core/src/main/java/com/tencent/trpc/core/filter/ConsumerInvokerHeadFilter.java b/trpc-core/src/main/java/com/tencent/trpc/core/filter/ConsumerInvokerHeadFilter.java
@@ -1,7 +1,7 @@
 /*
  * Tencent is pleased to support the open source community by making tRPC available.
  *
- * Copyright (C) 2023 THL A29 Limited, a Tencent company. 
+ * Copyright (C) 2023 THL A29 Limited, a Tencent company.
  * All rights reserved.
  *
  * If you have downloaded a copy of the tRPC source code from Tencent,
@@ -20,6 +20,10 @@
 import com.tencent.trpc.core.rpc.RequestMeta;
 import com.tencent.trpc.core.rpc.Response;
 import com.tencent.trpc.core.rpc.RpcContext;
+import com.tencent.trpc.core.rpc.RpcContextValueKeys;
+import com.tencent.trpc.core.utils.RpcContextUtils;
+
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 
@@ -43,6 +47,7 @@ public CompletionStage<Response> filter(Invoker<?> invoker, Request request) {
         ConsumerInvoker consumerInvoker = (ConsumerInvoker) invoker;
         // combine with DefClusterInvocationHandler#genRequest to complete the information of the request.
         prepareRequestInfoBeforeInvoke(request, consumerInvoker);
+        contextWithRemoteCalleeAddr(context, request);
         startLog(context, request);
         CompletableFuture<Response> future = invoker.invoke(request).toCompletableFuture();
         if (logger.isDebugEnabled()) {
@@ -60,6 +65,24 @@ public CompletionStage<Response> filter(Invoker<?> invoker, Request request) {
         });
     }
 
+    /**
+     * Set the request remote callee addr to RpcContext.
+     * Caller IP: CTX_CALLEE_REMOTE_IP
+     * Caller PORT: CTX_CALLEE_REMOTE_PORT
+     *
+     * @param context RpcContext
+     * @param request Request
+     */
+    private void contextWithRemoteCalleeAddr(RpcContext context, Request request) {
+        Optional.ofNullable(request.getMeta().getRemoteAddress()).ifPresent(remoteAddr
+                -> {
+            RpcContextUtils.putValueMapValue(context, RpcContextValueKeys.CTX_CALLEE_REMOTE_IP,
+                    remoteAddr.getAddress().getHostAddress());
+            RpcContextUtils.putValueMapValue(context, RpcContextValueKeys.CTX_CALLEE_REMOTE_PORT,
+                    remoteAddr.getPort());
+        });
+    }
+
     private void prepareRequestInfoBeforeInvoke(Request request,
             ConsumerInvoker<?> consumerInvoker) {
         RequestMeta meta = request.getMeta();

diff --git a/trpc-core/src/main/java/com/tencent/trpc/core/filter/ProviderInvokerHeadFilter.java b/trpc-core/src/main/java/com/tencent/trpc/core/filter/ProviderInvokerHeadFilter.java
@@ -52,7 +52,7 @@ public int getOrder() {
     public CompletionStage<Response> filter(Invoker<?> invoker, Request request) {
         RpcServerContext serverContext = (RpcServerContext) (request.getContext());
         prepareRequestInfoBeforeInvoke(request, (ProviderInvoker) invoker);
-        contextWithRemoteIp(serverContext, request);
+        contextWithRemoteCallerAddr(serverContext, request);
         startLog(serverContext, request);
         CompletableFuture<Response> future = invoker.invoke(request).toCompletableFuture();
         if (logger.isDebugEnabled()) {
@@ -62,15 +62,21 @@ public CompletionStage<Response> filter(Invoker<?> invoker, Request request) {
     }
 
     /**
-     * Set the request remote IP to RpcServerContext, with the key as CTX_CALLER_REMOTE_IP.
+     * Set the request remote caller addr to RpcServerContext.
+     * Caller IP: CTX_CALLER_REMOTE_IP
+     * Caller PORT: CTX_CALLER_REMOTE_PORT
      *
      * @param serverContext RpcServerContext
      * @param request Request
      */
-    private void contextWithRemoteIp(RpcServerContext serverContext, Request request) {
+    private void contextWithRemoteCallerAddr(RpcServerContext serverContext, Request request) {
         Optional.ofNullable(request.getMeta().getRemoteAddress()).ifPresent(remoteAddr
-                -> RpcContextUtils.putValueMapValue(serverContext, RpcContextValueKeys.CTX_CALLER_REMOTE_IP,
-                remoteAddr.getAddress().getHostAddress()));
+                -> {
+            RpcContextUtils.putValueMapValue(serverContext, RpcContextValueKeys.CTX_CALLER_REMOTE_IP,
+                    remoteAddr.getAddress().getHostAddress());
+            RpcContextUtils.putValueMapValue(serverContext, RpcContextValueKeys.CTX_CALLER_REMOTE_PORT,
+                    remoteAddr.getPort());
+        });
 
     }
 

diff --git a/trpc-core/src/main/java/com/tencent/trpc/core/rpc/RpcContextValueKeys.java b/trpc-core/src/main/java/com/tencent/trpc/core/rpc/RpcContextValueKeys.java
@@ -24,9 +24,21 @@ public class RpcContextValueKeys {
      */
     public static final String CTX_TELEMETRY_TRACE_SPAN = "ctx_telemetry_trace_span";
     /**
-     * Remote IP.
+     * Caller Remote IP.
      */
     public static final String CTX_CALLER_REMOTE_IP = "ctx_caller_remote_ip";
+    /**
+     * Caller Remote PORT.
+     */
+    public static final String CTX_CALLER_REMOTE_PORT = "ctx_caller_remote_port";
+    /**
+     * Callee Remote IP.
+     */
+    public static final String CTX_CALLEE_REMOTE_IP = "ctx_callee_remote_ip";
+    /**
+     * Callee Remote PORT.
+     */
+    public static final String CTX_CALLEE_REMOTE_PORT = "ctx_callee_remote_port";
     /**
      * Full link timeout key.
      */

diff --git a/...c-proto-standard/src/test/java/com/tencent/trpc/proto/standard/common/TRpcServerTest.java b/...c-proto-standard/src/test/java/com/tencent/trpc/proto/standard/common/TRpcServerTest.java
@@ -1,7 +1,7 @@
 /*
  * Tencent is pleased to support the open source community by making tRPC available.
  *
- * Copyright (C) 2023 THL A29 Limited, a Tencent company. 
+ * Copyright (C) 2023 THL A29 Limited, a Tencent company.
  * All rights reserved.
  *
  * If you have downloaded a copy of the tRPC source code from Tencent,
@@ -29,6 +29,7 @@
 import com.tencent.trpc.core.exception.TransportException;
 import com.tencent.trpc.core.rpc.RpcClient;
 import com.tencent.trpc.core.rpc.RpcClientContext;
+import com.tencent.trpc.core.rpc.RpcContextValueKeys;
 import com.tencent.trpc.core.rpc.RpcServer;
 import com.tencent.trpc.core.rpc.RpcServerManager;
 import com.tencent.trpc.core.transport.Channel;
@@ -241,6 +242,7 @@ public void serverNormalTest() {
                 .println(">>>>>>>>>" + new String((byte[]) (context.getRspAttachMap().get("key"))));
         assertEquals(new String((byte[]) (context.getRspAttachMap().get("key")), Charsets.UTF_8),
                 "abc");
+        assertEquals("127.0.0.1",context.getValueMap().get(RpcContextValueKeys.CTX_CALLEE_REMOTE_IP));
     }
 
     @Test