From 59a215fc4452cd5db576e046e2c5f639da61f21c Mon Sep 17 00:00:00 2001
From: gongdewei <kylixs@qq.com>
Date: Tue, 25 May 2021 14:33:27 +0800
Subject: [PATCH] support set agent id with request

---
 .../server/web/LegacyApiController.java       | 12 +++++++++--
 .../shell/term/impl/http/api/ApiRequest.java  | 20 ++++++++++++++-----
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/channel/channel-server/src/main/java/com/alibaba/arthas/channel/server/web/LegacyApiController.java b/channel/channel-server/src/main/java/com/alibaba/arthas/channel/server/web/LegacyApiController.java
index d4a67e6f4..2371075fb 100644
--- a/channel/channel-server/src/main/java/com/alibaba/arthas/channel/server/web/LegacyApiController.java
+++ b/channel/channel-server/src/main/java/com/alibaba/arthas/channel/server/web/LegacyApiController.java
@@ -49,11 +49,19 @@ public class LegacyApiController {
     @Autowired
     private ApiActionDelegateService apiActionDelegateService;
 
-    @RequestMapping("/legacy_api/{agentId}")
-    public Mono<ApiResponse> process(@PathVariable String agentId, @RequestBody ApiRequest request) {
+    @RequestMapping({"/legacy_api/{agentId}", "/legacy_api"})
+    public Mono<ApiResponse> process(@PathVariable(required = false) String agentId, @RequestBody ApiRequest request) {
 
         try {
+            // check agentId
+            if (StringUtils.hasText(request.getAgentId())) {
+                if (agentId != null && !agentId.equals(request.getAgentId())) {
+                    throw new IllegalArgumentException("Inconsistent agentId in path and request: " + agentId + ", " + request.getAgentId());
+                }
+                agentId = request.getAgentId();
+            }
             checkAgentExists(agentId);
+            request.setAgentId(agentId);
 
             // set default exec timeout
             if (request.getExecTimeout() == null) {
diff --git a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/api/ApiRequest.java b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/api/ApiRequest.java
index e57a68fd2..353024bbb 100644
--- a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/api/ApiRequest.java
+++ b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/api/ApiRequest.java
@@ -10,6 +10,7 @@ public class ApiRequest {
     private String action;
     private String command;
     private String requestId;
+    private String agentId;
     private String sessionId;
     private String consumerId;
     private Integer execTimeout;
@@ -18,11 +19,12 @@ public class ApiRequest {
     public String toString() {
         return "ApiRequest{" +
                 "action='" + action + '\'' +
-                ", command='" + command + '\'' +
-                ", requestId='" + requestId + '\'' +
-                ", sessionId='" + sessionId + '\'' +
-                ", consumerId='" + consumerId + '\'' +
-                ", execTimeout=" + execTimeout +
+                (command != null ? ", command='" + command + '\'' : "") +
+                (requestId != null ? ", requestId='" + requestId + '\'' : "") +
+                (agentId != null ? ", agentId='" + agentId + '\'' : "") +
+                (sessionId != null ? ", sessionId='" + sessionId + '\'' : "") +
+                (consumerId != null ? ", consumerId='" + consumerId + '\'' : "") +
+                (execTimeout != null ? ", execTimeout=" + execTimeout : "") +
                 '}';
     }
 
@@ -50,6 +52,14 @@ public class ApiRequest {
         this.requestId = requestId;
     }
 
+    public String getAgentId() {
+        return agentId;
+    }
+
+    public void setAgentId(String agentId) {
+        this.agentId = agentId;
+    }
+
     public String getSessionId() {
         return sessionId;
     }