From 6aee49083e67d33a2d79d999f42f288c65c3ca4c Mon Sep 17 00:00:00 2001
From: Andrew Azores <me@andrewazor.es>
Date: Fri, 1 Nov 2019 17:45:29 -0400
Subject: [PATCH] Improve handling of JMXServiceURLs

---
 build.gradle                                  |  2 +-
 .../containerjfr/core/net/JFRConnection.java  | 42 +++++++++----------
 .../core/net/JFRConnectionToolkit.java        |  6 ++-
 3 files changed, 25 insertions(+), 25 deletions(-)

diff --git a/build.gradle b/build.gradle
index 3980fcbe..93eaeed0 100644
--- a/build.gradle
+++ b/build.gradle
@@ -42,7 +42,7 @@ dependencies {
 }
 
 publishing {
-    def coreVersion = '0.4.0';
+    def coreVersion = '0.4.1';
     publications {
         maven(MavenPublication) {
             groupId = 'com.redhat.rhjmc'
diff --git a/src/main/java/com/redhat/rhjmc/containerjfr/core/net/JFRConnection.java b/src/main/java/com/redhat/rhjmc/containerjfr/core/net/JFRConnection.java
index d4c45de2..4821d303 100644
--- a/src/main/java/com/redhat/rhjmc/containerjfr/core/net/JFRConnection.java
+++ b/src/main/java/com/redhat/rhjmc/containerjfr/core/net/JFRConnection.java
@@ -1,39 +1,32 @@
 package com.redhat.rhjmc.containerjfr.core.net;
 
-import javax.management.remote.JMXServiceURL;
+import java.io.IOException;
 
 import com.redhat.rhjmc.containerjfr.core.sys.Clock;
 import com.redhat.rhjmc.containerjfr.core.tui.ClientWriter;
 
+import org.openjdk.jmc.rjmx.ConnectionToolkit;
+import org.openjdk.jmc.rjmx.IConnectionDescriptor;
 import org.openjdk.jmc.rjmx.IConnectionHandle;
 import org.openjdk.jmc.rjmx.IConnectionListener;
 import org.openjdk.jmc.rjmx.internal.DefaultConnectionHandle;
-import org.openjdk.jmc.rjmx.internal.JMXConnectionDescriptor;
 import org.openjdk.jmc.rjmx.internal.RJMXConnection;
 import org.openjdk.jmc.rjmx.internal.ServerDescriptor;
 import org.openjdk.jmc.rjmx.services.jfr.IFlightRecorderService;
 import org.openjdk.jmc.rjmx.services.jfr.internal.FlightRecorderServiceFactory;
-import org.openjdk.jmc.ui.common.security.InMemoryCredentials;
 
 public class JFRConnection implements AutoCloseable {
 
-    static final String URL_FORMAT = "service:jmx:rmi:///jndi/rmi://%s:%d/jmxrmi";
     public static final int DEFAULT_PORT = 9091;
 
     private final ClientWriter cw;
-    private final JMXServiceURL url;
     private final RJMXConnection rjmxConnection;
     private final IConnectionHandle handle;
     private final IFlightRecorderService service;
 
-    JFRConnection(ClientWriter cw, String host, int port) throws Exception {
-        this(cw, new JMXServiceURL(String.format(URL_FORMAT, host, port)));
-    }
-
-    JFRConnection(ClientWriter cw, JMXServiceURL url) throws Exception {
+    JFRConnection(ClientWriter cw, IConnectionDescriptor cd) throws Exception {
         this.cw = cw;
-        this.url = url;
-        this.rjmxConnection = attemptConnect(url);
+        this.rjmxConnection = attemptConnect(cd);
         this.handle = new DefaultConnectionHandle(rjmxConnection, "RJMX Connection", new IConnectionListener[0]);
         this.service = new FlightRecorderServiceFactory().getServiceInstance(handle);
     }
@@ -47,15 +40,25 @@ public IFlightRecorderService getService() {
     }
 
     public long getApproximateServerTime(Clock clock) {
-        return rjmxConnection.getApproximateServerTime(clock.getWallTime());
+        return this.rjmxConnection.getApproximateServerTime(clock.getWallTime());
     }
 
     public String getHost() {
-        return this.url.getHost();
+        try {
+            return ConnectionToolkit.getHostName(this.rjmxConnection.getConnectionDescriptor().createJMXServiceURL());
+        } catch (IOException e) {
+            cw.println(e);
+            return "unknown";
+        }
     }
 
     public int getPort() {
-        return this.url.getPort();
+        try {
+            return ConnectionToolkit.getPort(this.rjmxConnection.getConnectionDescriptor().createJMXServiceURL());
+        } catch (IOException e) {
+            cw.println(e);
+            return 0;
+        }
     }
 
     public void disconnect() {
@@ -67,14 +70,9 @@ public void close() {
         this.disconnect();
     }
 
-    private RJMXConnection attemptConnect(JMXServiceURL url) throws Exception {
-        JMXConnectionDescriptor cd = new JMXConnectionDescriptor(
-                url,
-                new InMemoryCredentials(null, null));
-        ServerDescriptor sd = new ServerDescriptor(null, "Container", null);
-
+    private RJMXConnection attemptConnect(IConnectionDescriptor cd) throws Exception {
         try {
-            RJMXConnection conn = new RJMXConnection(cd, sd, JFRConnection::failConnection);
+            RJMXConnection conn = new RJMXConnection(cd, new ServerDescriptor(), JFRConnection::failConnection);
             if (!conn.connect()) {
                 failConnection();
             }
diff --git a/src/main/java/com/redhat/rhjmc/containerjfr/core/net/JFRConnectionToolkit.java b/src/main/java/com/redhat/rhjmc/containerjfr/core/net/JFRConnectionToolkit.java
index 116befd5..fddcf262 100644
--- a/src/main/java/com/redhat/rhjmc/containerjfr/core/net/JFRConnectionToolkit.java
+++ b/src/main/java/com/redhat/rhjmc/containerjfr/core/net/JFRConnectionToolkit.java
@@ -4,6 +4,8 @@
 
 import com.redhat.rhjmc.containerjfr.core.tui.ClientWriter;
 
+import org.openjdk.jmc.rjmx.ConnectionDescriptorBuilder;
+
 public class JFRConnectionToolkit {
 
     private final ClientWriter cw;
@@ -13,7 +15,7 @@ public JFRConnectionToolkit(ClientWriter cw) {
     }
 
     public JFRConnection connect(JMXServiceURL url) throws Exception {
-        return new JFRConnection(cw, url);
+        return new JFRConnection(cw, new ConnectionDescriptorBuilder().url(url).build());
     }
 
     public JFRConnection connect(String host) throws Exception {
@@ -21,6 +23,6 @@ public JFRConnection connect(String host) throws Exception {
     }
 
     public JFRConnection connect(String host, int port) throws Exception {
-        return new JFRConnection(cw, host, port);
+        return new JFRConnection(cw, new ConnectionDescriptorBuilder().hostName(host).port(port).build());
     }
 }