From 693cc8d0adba2426b66d1e40a5ff524fd701ad40 Mon Sep 17 00:00:00 2001
From: Haotian Zhang <928016560@qq.com>
Date: Fri, 4 Aug 2023 16:00:39 +0800
Subject: [PATCH] fix:fix consul token read null bug.

---
 .../client/flow/DefaultDiscoveryFlow.java     | 23 +++++++++---------
 .../connector/consul/ConsulAPIConnector.java  | 24 +++++++++++++------
 .../connector/consul/ConsulContext.java       | 10 ++++++++
 3 files changed, 39 insertions(+), 18 deletions(-)

diff --git a/polaris-discovery/polaris-discovery-client/src/main/java/com/tencent/polaris/discovery/client/flow/DefaultDiscoveryFlow.java b/polaris-discovery/polaris-discovery-client/src/main/java/com/tencent/polaris/discovery/client/flow/DefaultDiscoveryFlow.java
index 13397c56e..fe0bd619d 100644
--- a/polaris-discovery/polaris-discovery-client/src/main/java/com/tencent/polaris/discovery/client/flow/DefaultDiscoveryFlow.java
+++ b/polaris-discovery/polaris-discovery-client/src/main/java/com/tencent/polaris/discovery/client/flow/DefaultDiscoveryFlow.java
@@ -49,9 +49,10 @@
 import com.tencent.polaris.client.api.ServiceCallResultListener;
 import com.tencent.polaris.client.util.Utils;
 import com.tencent.polaris.logging.LoggerFactory;
+import org.slf4j.Logger;
+
 import java.util.List;
 import java.util.Map;
-import org.slf4j.Logger;
 
 public class DefaultDiscoveryFlow implements DiscoveryFlow {
 
@@ -151,7 +152,7 @@ private InstanceRegisterResponse doRegister(InstanceRegisterRequest req, Map<Str
             CommonProviderRequest request = req.getRequest();
             try {
                 CommonProviderResponse response = serverConnector.registerInstance(request, customHeader);
-                LOG.info("register {}/{} instance {} succ", req.getNamespace(), req.getService(),
+                LOG.info("register {}/{} instance {} successfully", req.getNamespace(), req.getService(),
                         response.getInstanceID());
                 serviceCallResult.setRetStatus(RetStatus.RetSuccess);
                 serviceCallResult.setRetCode(ErrorCode.Success.getCode());
@@ -189,7 +190,7 @@ public void deRegister(InstanceDeregisterRequest req) {
                 serverConnector.deregisterInstance(request);
                 serviceCallResult.setRetStatus(RetStatus.RetSuccess);
                 serviceCallResult.setRetCode(ErrorCode.Success.getCode());
-                LOG.info("deregister instance {} succ", req);
+                LOG.info("deregister instance {} successfully", req);
                 return;
             } catch (PolarisException e) {
                 serviceCallResult.setRetStatus(RetStatus.RetFail);
@@ -288,8 +289,8 @@ private void reportInvokeStat(ServiceCallResult req) throws PolarisException {
      * 上报内部服务调用结果
      *
      * @param serviceCallResult 服务调用结果
-     * @param targetServer 目标服务端
-     * @param method 方法
+     * @param targetServer      目标服务端
+     * @param method            方法
      */
     public void reportServerCall(ServiceCallResult serviceCallResult, TargetServer targetServer, String method) {
         if (null != targetServer) {
@@ -303,10 +304,10 @@ public void reportServerCall(ServiceCallResult serviceCallResult, TargetServer t
         reportInvokeStat(serviceCallResult);
     }
 
-	@Override
-	public void destroy() {
-		if (registerFlow != null) {
-			registerFlow.destroy();
-		}
-	}
+    @Override
+    public void destroy() {
+        if (registerFlow != null) {
+            registerFlow.destroy();
+        }
+    }
 }
diff --git a/polaris-plugins/polaris-plugins-connector/connector-consul/src/main/java/com/tencent/polaris/plugins/connector/consul/ConsulAPIConnector.java b/polaris-plugins/polaris-plugins-connector/connector-consul/src/main/java/com/tencent/polaris/plugins/connector/consul/ConsulAPIConnector.java
index d0da65790..66cb849bd 100644
--- a/polaris-plugins/polaris-plugins-connector/connector-consul/src/main/java/com/tencent/polaris/plugins/connector/consul/ConsulAPIConnector.java
+++ b/polaris-plugins/polaris-plugins-connector/connector-consul/src/main/java/com/tencent/polaris/plugins/connector/consul/ConsulAPIConnector.java
@@ -32,6 +32,7 @@
 import com.ecwid.consul.v1.catalog.CatalogServicesRequest;
 import com.ecwid.consul.v1.health.HealthServicesRequest;
 import com.ecwid.consul.v1.health.model.HealthService;
+import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.tencent.polaris.api.config.global.ServerConnectorConfig;
 import com.tencent.polaris.api.config.plugin.DefaultPlugins;
@@ -62,7 +63,6 @@
 import com.tencent.polaris.plugins.connector.common.DestroyableServerConnector;
 import com.tencent.polaris.plugins.connector.common.ServiceInstancesResponse;
 import com.tencent.polaris.plugins.connector.common.ServiceUpdateTask;
-import com.tencent.polaris.plugins.connector.common.constant.ConsulConstant;
 import org.slf4j.Logger;
 
 import java.util.ArrayList;
@@ -148,8 +148,8 @@ public void init(InitContext ctx) throws PolarisException {
             if (CollectionUtils.isNotEmpty(serverConnectorConfigs)) {
                 for (ServerConnectorConfigImpl serverConnectorConfig : serverConnectorConfigs) {
                     if (DefaultPlugins.SERVER_CONNECTOR_CONSUL.equals(serverConnectorConfig.getProtocol())) {
-                        initActually(ctx, serverConnectorConfig);
                         mapper = new ObjectMapper();
+                        initActually(ctx, serverConnectorConfig);
                     }
                 }
             }
@@ -169,6 +169,7 @@ private void initActually(InitContext ctx, ServerConnectorConfig connectorConfig
         int lastIndex = address.lastIndexOf(":");
         String agentHost = address.substring(0, lastIndex);
         int agentPort = Integer.parseInt(address.substring(lastIndex + 1));
+        LOG.debug("Consul Server : [" + address + "]");
         consulRawClient = new ConsulRawClient(agentHost, agentPort);
         consulClient = new ConsulClient(consulRawClient);
 
@@ -226,8 +227,8 @@ public CommonProviderResponse registerInstance(CommonProviderRequest req, Map<St
                 NewService service = buildRegisterInstanceRequest(req);
                 String json = getGson().toJson(service);
                 HttpResponse rawResponse;
-                if (req.getMetadata().containsKey(ConsulConstant.MetadataMapKey.TOKEN_KEY)) {
-                    String token = req.getMetadata().get(ConsulConstant.MetadataMapKey.TOKEN_KEY);
+                if (StringUtils.isNotBlank(req.getToken())) {
+                    String token = req.getToken();
                     UrlParameters tokenParam = token != null ? new SingleUrlParameters("token", token) : null;
                     rawResponse = consulRawClient.makePutRequest("/v1/agent/service/register", json,
                             tokenParam);
@@ -235,10 +236,19 @@ public CommonProviderResponse registerInstance(CommonProviderRequest req, Map<St
                     rawResponse = consulRawClient.makePutRequest("/v1/agent/service/register", json);
                 }
                 if (rawResponse.getStatusCode() != 200) {
+                    try {
+                        LOG.warn(new ObjectMapper().writeValueAsString(rawResponse));
+                    } catch (JsonProcessingException ignore) {
+                    }
                     throw new OperationException(rawResponse);
                 }
                 CommonProviderResponse resp = new CommonProviderResponse();
                 consulContext.setInstanceId(service.getId());
+                String checkId = consulContext.getInstanceId();
+                if (!checkId.startsWith("service:")) {
+                    checkId = "service:" + checkId;
+                }
+                consulContext.setCheckId(checkId);
                 resp.setInstanceID(service.getId());
                 resp.setExists(false);
                 LOG.info("Registered service to Consul: " + service);
@@ -304,7 +314,7 @@ public void deregisterInstance(CommonProviderRequest req) throws PolarisExceptio
             ServiceKey serviceKey = new ServiceKey(req.getNamespace(), req.getService());
             try {
                 LOG.info("Unregistering service to Consul: " + consulContext.getInstanceId());
-                this.consulClient.agentServiceDeregister(consulContext.getInstanceId());
+                this.consulClient.agentServiceDeregister(consulContext.getInstanceId(), req.getToken());
                 LOG.info("Unregistered service to Consul: " + consulContext.getInstanceId());
                 ieRegistered = false;
             } catch (ConsulException e) {
@@ -320,9 +330,9 @@ public void heartbeat(CommonProviderRequest req) throws PolarisException {
         if (ieRegistered) {
             ServiceKey serviceKey = new ServiceKey(req.getNamespace(), req.getService());
             try {
-                this.consulClient.agentCheckPass("service:" + consulContext.getInstanceId());
+                this.consulClient.agentCheckPass(consulContext.getCheckId(), null, req.getToken());
                 if (LOG.isDebugEnabled()) {
-                    LOG.debug("Heartbeat service to Consul: " + consulContext.getInstanceId());
+                    LOG.debug("Heartbeat service to Consul: " + consulContext.getCheckId());
                 }
             } catch (ConsulException e) {
                 throw new RetriableException(ErrorCode.NETWORK_ERROR,
diff --git a/polaris-plugins/polaris-plugins-connector/connector-consul/src/main/java/com/tencent/polaris/plugins/connector/consul/ConsulContext.java b/polaris-plugins/polaris-plugins-connector/connector-consul/src/main/java/com/tencent/polaris/plugins/connector/consul/ConsulContext.java
index cfa5a8891..5d18368a3 100644
--- a/polaris-plugins/polaris-plugins-connector/connector-consul/src/main/java/com/tencent/polaris/plugins/connector/consul/ConsulContext.java
+++ b/polaris-plugins/polaris-plugins-connector/connector-consul/src/main/java/com/tencent/polaris/plugins/connector/consul/ConsulContext.java
@@ -44,6 +44,8 @@ public class ConsulContext {
 
     private NewService.Check check;
 
+    private String checkId;
+
     public ConsulContext() {
         serviceName = "";
         instanceId = "";
@@ -108,4 +110,12 @@ public NewService.Check getCheck() {
     public void setCheck(NewService.Check check) {
         this.check = check;
     }
+
+    public String getCheckId() {
+        return checkId;
+    }
+
+    public void setCheckId(String checkId) {
+        this.checkId = checkId;
+    }
 }