Merge pull request #14 from newrelic/preconnect-original

Always use configured host on preconnect
newrelic · Aug 20, 2020 · 010b9fb · 010b9fb
2 parents cd917af + 760ca95
commit 010b9fb
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 14 deletions.
diff --git a/newrelic-agent/src/main/java/com/newrelic/agent/transport/DataSenderImpl.java b/newrelic-agent/src/main/java/com/newrelic/agent/transport/DataSenderImpl.java
@@ -101,7 +101,8 @@ public class DataSenderImpl implements DataSender {
 
     private final HttpClientWrapper httpClientWrapper;
 
-    private volatile String host;
+    private final String originalHost;
+    private volatile String redirectHost;
     private final int port;
     private volatile String protocol;
     private volatile boolean auditMode;
@@ -129,7 +130,8 @@ public DataSenderImpl(
         this.logger = logger;
         this.configService = configService;
         logger.info(MessageFormat.format("Setting audit_mode to {0}", auditMode));
-        host = config.getHost();
+        originalHost = config.getHost();
+        redirectHost = config.getHost();
         port = config.getPort();
 
         protocol = config.isSSL() ? "https" : "http";
@@ -187,8 +189,8 @@ Object getAgentRunId() {
     public Map<String, Object> connect(Map<String, Object> startupOptions) throws Exception {
         String redirectHost = parsePreconnectAndReturnHost();
         if (redirectHost != null) {
-            host = redirectHost;
-            logger.info(MessageFormat.format("Collector redirection to {0}:{1}", host, Integer.toString(port)));
+            this.redirectHost = redirectHost;
+            logger.info(MessageFormat.format("Collector redirection to {0}:{1}", this.redirectHost, Integer.toString(port)));
         } else if (configService.getDefaultAgentConfig().laspEnabled()) {
             throw new ForceDisconnectException("The agent did not receive one or more security policies that it expected and will shut down."
                     + " Please contact support.");
@@ -207,7 +209,7 @@ private String parsePreconnectAndReturnHost() throws Exception {
         }
         token.put("high_security", agentConfig.isHighSecurity());
         params.add(token);
-        Object response = invokeNoRunId(CollectorMethods.PRECONNECT, compressedEncoding, params);
+        Object response = invokeNoRunId(originalHost, CollectorMethods.PRECONNECT, compressedEncoding, params);
 
         if (response != null) {
             Map<?, ?> returnValue = (Map<?, ?>) response;
@@ -232,7 +234,7 @@ private Map<String, Object> doConnect(Map<String, Object> startupOptions) throws
         startupOptions.put(ENV_METADATA, metadata);
 
         params.add(startupOptions);
-        Object response = invokeNoRunId(CollectorMethods.CONNECT, compressedEncoding, params);
+        Object response = invokeNoRunId(redirectHost, CollectorMethods.CONNECT, compressedEncoding, params);
         if (!(response instanceof Map)) {
             throw new UnexpectedException(MessageFormat.format("Expected a map of connection data, got {0}", response));
         }
@@ -487,17 +489,17 @@ void setMaxPayloadSizeInBytes(int payloadSizeInBytes) {
 
     private Object invokeRunId(String method, String encoding, Object runId, JSONStreamAware params) throws Exception {
         String uri = MessageFormat.format(agentRunIdUriPattern, method, runId.toString());
-        return invoke(method, encoding, uri, params);
+        return invoke(redirectHost, method, encoding, uri, params);
     }
 
-    private Object invokeNoRunId(String method, String encoding, JSONStreamAware params) throws Exception {
+    private Object invokeNoRunId(String host, String method, String encoding, JSONStreamAware params) throws Exception {
         String uri = MessageFormat.format(noAgentRunIdUriPattern, method);
-        return invoke(method, encoding, uri, params);
+        return invoke(host, method, encoding, uri, params);
     }
 
-    private Object invoke(String method, String encoding, String uri, JSONStreamAware params) throws Exception {
+    private Object invoke(String host, String method, String encoding, String uri, JSONStreamAware params) throws Exception {
         // ReadResult should be from a valid 2xx response at this point otherwise send method throws an exception here
-        ReadResult readResult = send(method, encoding, uri, params);
+        ReadResult readResult = send(host, method, encoding, uri, params);
         Map<?, ?> responseMap = null;
         String responseBody = readResult.getResponseBody();
 
@@ -535,7 +537,7 @@ private Object invoke(String method, String encoding, String uri, JSONStreamAwar
      * response code value. The previous behavior of a 200 (“OK”) with an exact string in the body that should be
      * matched/parsed has been deprecated.
      */
-    private ReadResult connectAndSend(String method, String encoding, String uri, JSONStreamAware params) throws Exception {
+    private ReadResult connectAndSend(String host, String method, String encoding, String uri, JSONStreamAware params) throws Exception {
         byte[] data = writeData(encoding, params);
 
         /*
@@ -636,9 +638,9 @@ private boolean methodShouldBeAudited(String method) {
         return true;
     }
 
-    private ReadResult send(String method, String encoding, String uri, JSONStreamAware params) throws Exception {
+    private ReadResult send(String host, String method, String encoding, String uri, JSONStreamAware params) throws Exception {
         try {
-            return connectAndSend(method, encoding, uri, params);
+            return connectAndSend(host, method, encoding, uri, params);
         } catch (MalformedURLException e) {
             logger.log(Level.SEVERE, "You have requested a connection to New Relic via a protocol which is unavailable in your runtime: {0}", e.toString());
             throw new ForceDisconnectException(e.toString());

diff --git a/...elic-agent/src/test/java/com/newrelic/agent/transport/DataSenderImplConnectCycleTest.java b/...elic-agent/src/test/java/com/newrelic/agent/transport/DataSenderImplConnectCycleTest.java
@@ -64,6 +64,35 @@ public void testFullPreconnectConnectCycleNoSecurity() throws Exception {
         assertEquals("value", result.get("other"));
     }
 
+    @Test
+    public void testReconnectCycleNoSecurity() throws Exception {
+        MockitoAnnotations.initMocks(this);
+
+        final Map<String, Object> settings = new HashMap<>();
+        settings.put(AgentConfigImpl.APP_NAME, "Unit Test");
+        settings.put(AgentConfigImpl.HOST, "no-collector.example.com");
+
+        final MockServiceManager serviceManager = new MockServiceManager(
+                ConfigServiceFactory.createConfigServiceUsingSettings(settings)
+        );
+        serviceManager.setStatsService(mockStatsService);
+
+        HttpClientWrapper wrapper = new ConnectCycleNoSecuritySuccessClientWrapper();
+
+        DataSenderImpl target = new DataSenderImpl(serviceManager.getConfigService().getDefaultAgentConfig(), wrapper, null, Agent.LOG,
+                ServiceFactory.getConfigService());
+        target.setAgentRunId("agent run id");
+
+        Map<String, Object> startupOptions = new HashMap<>();
+        startupOptions.put("test-sentinel", "test-value");
+        Map<String, Object> result = target.connect(startupOptions);
+        assertEquals("my-run-id", target.getAgentRunId());
+        assertEquals("value", result.get("other"));
+
+        // This will assert that preconnect was called with the original host, not the redirect host.
+        target.connect(startupOptions);
+    }
+
     @Test
     public void testFullPreconnectConnectCycleWithLasp() throws Exception {
         MockitoAnnotations.initMocks(this);