Getting authentication credential prior to the socket opening

dzmipt · Nov 28, 2023 · 14ee348 · 14ee348
1 parent 0ae88ed
commit 14ee348
Show file tree

Hide file tree

Showing 6 changed files with 49 additions and 55 deletions.
diff --git a/src/kx/KAuthentication.java b/src/kx/KAuthentication.java
@@ -0,0 +1,8 @@
+package kx;
+
+import java.io.IOException;
+
+public interface KAuthentication {
+
+    String getUserPassword() throws IOException;
+}
diff --git a/src/kx/KConnection.java b/src/kx/KConnection.java
@@ -12,7 +12,6 @@ public class KConnection {
 
     private final String host;
     private final int port;
-    private final String userPassword;
     private final boolean useTLS;
     private volatile boolean closed = true;
 
@@ -22,6 +21,7 @@ public class KConnection {
 
     private SocketReader socketReader;
     private ConnectionStateListener connectionStateListener = null;
+    private KAuthentication authentication = null;
 
     void io(Socket s) throws IOException {
         s.setTcpNoDelay(true);
@@ -62,6 +62,7 @@ public boolean isClosed() {
     }
 
     private void connect() throws IOException, K4AccessException {
+        String userPassword = authentication == null ? "" : authentication.getUserPassword();
         s = new Socket();
         s.setReceiveBufferSize(1024 * 1024);
         s.connect(new InetSocketAddress(host, port));
@@ -96,10 +97,14 @@ private void connect() throws IOException, K4AccessException {
         }
     }
 
-    public KConnection(String h, int p, String userPassword, boolean useTLS) {
+    public KConnection(String h, int p, boolean useTLS) {
+        this(h, p, useTLS, null);
+    }
+
+    public KConnection(String h, int p, boolean useTLS, KAuthentication authentication) {
         host = h;
         port = p;
-        this.userPassword = userPassword;
+        this.authentication = authentication;
         this.useTLS = useTLS;
     }
 

diff --git a/src/studio/kdb/Session.java b/src/studio/kdb/Session.java
@@ -1,9 +1,6 @@
 package studio.kdb;
 
-import kx.ConnectionStateListener;
-import kx.K4Exception;
-import kx.KConnection;
-import kx.ProgressCallback;
+import kx.*;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import studio.core.AuthenticationManager;
@@ -17,13 +14,11 @@
 import java.util.List;
 import java.util.Map;
 
-public class Session implements ConnectionStateListener {
+public class Session implements ConnectionStateListener, KAuthentication {
     private KConnection kConn;
     private long created;
     private final Server server;
 
-    private boolean busy = false;
-
     private static final long HOUR = 1000*3600;
     private static SessionCreator sessionCreator = new SessionCreator();
     private final List<EditorTab> editors = new ArrayList<>();
@@ -69,8 +64,7 @@ private Session(Server server) {
     }
 
     private void init() {
-        log.info("Connecting to server " + server.getDescription(true));
-        kConn = createConnection(server);
+        kConn = sessionCreator.createConnection(server, this);
         if (kConn == null) throw new RuntimeException("Failure in the authentication plugin");
         created = System.currentTimeMillis();
         kConn.setConnectionStateListener(this);
@@ -81,22 +75,6 @@ static void mock(SessionCreator sessionCreator) {
         Session.sessionCreator = sessionCreator;
     }
 
-    private static KConnection createConnection(Server s) {
-        return sessionCreator.createConnection(s);
-    }
-
-    public boolean isBusy() {
-        return busy;
-    }
-
-    public void setFree() {
-        busy = false;
-    }
-
-    public void setBusy() {
-        busy = true;
-    }
-
     public K.KBase execute(K.KBase x, ProgressCallback progress) throws K4Exception, IOException, InterruptedException {
         return kConn.k(x, progress);
     }
@@ -124,27 +102,33 @@ public void validate() {
         }
     }
 
-    public static class SessionCreator {
-        public KConnection createConnection(Server s) {
-            try {
-                Class<?> clazz = AuthenticationManager.getInstance().lookup(s.getAuthenticationMechanism());
-                IAuthenticationMechanism authenticationMechanism = (IAuthenticationMechanism) clazz.newInstance();
-
-                authenticationMechanism.setProperties(s.getAsProperties());
-                Credentials credentials = authenticationMechanism.getCredentials();
-
-                KConnection kConn;
-                if (credentials.getUsername().length() > 0) {
-                    String p = credentials.getPassword();
-                    kConn = new KConnection(s.getHost(), s.getPort(), credentials.getUsername() + ((p.length() == 0) ? "" : ":" + p), s.getUseTLS());
-                } else {
-                    kConn = new KConnection(s.getHost(), s.getPort(), "", s.getUseTLS());
-                }
-                return kConn;
-            } catch (InstantiationException | IllegalAccessException | IllegalArgumentException ex) {
-                log.error("Failed to initialize connection", ex);
-                return null;
+    @Override
+    public String getUserPassword() throws IOException {
+        try {
+            log.info("Getting authentication credential for {} with auth.method {}",
+                    server.getDescription(true), server.getAuthenticationMechanism());
+
+            Class<?> clazz = AuthenticationManager.getInstance().lookup(server.getAuthenticationMechanism());
+            IAuthenticationMechanism authenticationMechanism = (IAuthenticationMechanism) clazz.newInstance();
+
+            authenticationMechanism.setProperties(server.getAsProperties());
+            Credentials credentials = authenticationMechanism.getCredentials();
+
+            if (credentials.getUsername().length() > 0) {
+                String p = credentials.getPassword();
+                return credentials.getUsername() + ((p.length() == 0) ? "" : ":" + p);
+            } else {
+                return "";
             }
+        } catch (InstantiationException | IllegalAccessException | IllegalArgumentException ex) {
+            log.error("Failed to initialize connection", ex);
+            throw new IOException("Failed to get credentials", ex);
+        }
+    }
+
+    public static class SessionCreator {
+        public KConnection createConnection(Server s, KAuthentication authentication) {
+            return new KConnection(s.getHost(), s.getPort(), s.getUseTLS(), authentication);
         }
     }
 

diff --git a/src/studio/ui/action/QueryExecutor.java b/src/studio/ui/action/QueryExecutor.java
@@ -138,10 +138,6 @@ protected QueryResult doInBackground() {
                     closeConnection();
                 }
                 result.setError(e);
-            } finally {
-                if (session != null) {
-                    session.setFree();
-                }
             }
             result.setExecutionTime(System.currentTimeMillis() - startTime);
             if (result.getError() != null) {

diff --git a/test-integration/studio/kdb/MockQSession.java b/test-integration/studio/kdb/MockQSession.java
@@ -1,6 +1,7 @@
 package studio.kdb;
 
 import kx.K4Exception;
+import kx.KAuthentication;
 import kx.KConnection;
 import kx.ProgressCallback;
 import org.apache.logging.log4j.LogManager;
@@ -73,7 +74,7 @@ public static void unlockResponse() {
     private boolean closed = true;
 
     MockQSession() {
-        super("no host", 0, "mock user", false);
+        super("no host", 0, false);
         index = sessionIndex++;
     }
 
@@ -150,7 +151,7 @@ public static class MockQSessionCreator extends Session.SessionCreator {
         private List<MockQSession> sessions = new ArrayList<>();
 
         @Override
-        public KConnection createConnection(Server server) {
+        public KConnection createConnection(Server server, KAuthentication authentication) {
             MockQSession session = new MockQSession();
             sessions.add(session);
             return session;

diff --git a/test/studio/kdb/KSerialiseTest.java b/test/studio/kdb/KSerialiseTest.java
@@ -21,7 +21,7 @@ public class KSerialiseTest {
 
     @BeforeAll
     public static void connect() throws K4Exception, IOException, InterruptedException {
-        kConn = new KConnection("localhost", Integer.parseInt(System.getenv("qTestPort")), "", false);
+        kConn = new KConnection("localhost", Integer.parseInt(System.getenv("qTestPort")),false);
         kConn.k(new K.KCharacterVector(".z.pg:{$[x~\"reset\";`.z.pg set value;x]}"));
     }