From f9bb25fdb28cc3a9d890db8ff4c30986bc3ac92e Mon Sep 17 00:00:00 2001
From: Michael Jumper <mike.jumper@guac-dev.org>
Date: Tue, 17 Mar 2015 12:11:29 -0700
Subject: [PATCH] GUAC-1132: Associate active socket with connection record.

---
 .../AbstractGuacamoleSocketService.java       | 23 +++++++++++--------
 .../jdbc/socket/ActiveConnectionRecord.java   | 19 +++++++++++++--
 2 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/extensions/guacamole-auth-jdbc/modules/guacamole-auth-jdbc-base/src/main/java/org/glyptodon/guacamole/auth/jdbc/socket/AbstractGuacamoleSocketService.java b/extensions/guacamole-auth-jdbc/modules/guacamole-auth-jdbc-base/src/main/java/org/glyptodon/guacamole/auth/jdbc/socket/AbstractGuacamoleSocketService.java
index 279363b7d..7276fbe95 100644
--- a/extensions/guacamole-auth-jdbc/modules/guacamole-auth-jdbc-base/src/main/java/org/glyptodon/guacamole/auth/jdbc/socket/AbstractGuacamoleSocketService.java
+++ b/extensions/guacamole-auth-jdbc/modules/guacamole-auth-jdbc-base/src/main/java/org/glyptodon/guacamole/auth/jdbc/socket/AbstractGuacamoleSocketService.java
@@ -352,11 +352,8 @@ public abstract class AbstractGuacamoleSocketService implements GuacamoleSocketS
      * The connection will be automatically released when it closes, or if it
      * fails to establish entirely.
      *
-     * @param user
-     *     The user for whom the connection is being established.
-     *
-     * @param connection
-     *     The connection the user is connecting to.
+     * @param activeConnection
+     *     The active connection record of the connection in use.
      *
      * @param info
      *     Information describing the Guacamole client connecting to the given
@@ -370,7 +367,7 @@ public abstract class AbstractGuacamoleSocketService implements GuacamoleSocketS
      *     If an error occurs while the connection is being established, or
      *     while connection configuration information is being retrieved.
      */
-    private GuacamoleSocket getGuacamoleSocket(ActiveConnectionRecord activeConnection,
+    private GuacamoleSocket assignGuacamoleSocket(ActiveConnectionRecord activeConnection,
             GuacamoleClientInformation info)
             throws GuacamoleException {
 
@@ -382,13 +379,19 @@ public abstract class AbstractGuacamoleSocketService implements GuacamoleSocketS
         activeConnections.put(connection.getIdentifier(), activeConnection);
         activeConnectionGroups.put(connection.getParentIdentifier(), activeConnection);
 
-        // Return new socket
         try {
-            return new ConfiguredGuacamoleSocket(
+
+            // Obtain socket which will automatically run the cleanup task
+            GuacamoleSocket socket = new ConfiguredGuacamoleSocket(
                 getUnconfiguredGuacamoleSocket(cleanupTask),
                 getGuacamoleConfiguration(activeConnection.getUser(), connection),
                 info
             );
+
+            // Assign and return new socket
+            activeConnection.setActiveSocket(socket);
+            return socket;
+            
         }
 
         // Execute cleanup if socket could not be created
@@ -462,7 +465,7 @@ public abstract class AbstractGuacamoleSocketService implements GuacamoleSocketS
 
         // Acquire and connect to single connection
         acquire(user, Collections.singletonList(connection));
-        return getGuacamoleSocket(new ActiveConnectionRecord(user, connection), info);
+        return assignGuacamoleSocket(new ActiveConnectionRecord(user, connection), info);
 
     }
 
@@ -487,7 +490,7 @@ public abstract class AbstractGuacamoleSocketService implements GuacamoleSocketS
 
         // Acquire and connect to any child
         ModeledConnection connection = acquire(user, connections);
-        return getGuacamoleSocket(new ActiveConnectionRecord(user, connectionGroup, connection), info);
+        return assignGuacamoleSocket(new ActiveConnectionRecord(user, connectionGroup, connection), info);
 
     }
 
diff --git a/extensions/guacamole-auth-jdbc/modules/guacamole-auth-jdbc-base/src/main/java/org/glyptodon/guacamole/auth/jdbc/socket/ActiveConnectionRecord.java b/extensions/guacamole-auth-jdbc/modules/guacamole-auth-jdbc-base/src/main/java/org/glyptodon/guacamole/auth/jdbc/socket/ActiveConnectionRecord.java
index c7c46470d..708148848 100644
--- a/extensions/guacamole-auth-jdbc/modules/guacamole-auth-jdbc-base/src/main/java/org/glyptodon/guacamole/auth/jdbc/socket/ActiveConnectionRecord.java
+++ b/extensions/guacamole-auth-jdbc/modules/guacamole-auth-jdbc-base/src/main/java/org/glyptodon/guacamole/auth/jdbc/socket/ActiveConnectionRecord.java
@@ -62,6 +62,12 @@ public class ActiveConnectionRecord implements ConnectionRecord {
      */
     private final Date startDate = new Date();
 
+    /**
+     * The GuacamoleSocket used by the connection associated with this
+     * connection record.
+     */
+    private GuacamoleSocket socket;
+    
     /**
      * Creates a new connection record associated with the given user,
      * connection, and balancing connection group. The given balancing
@@ -189,8 +195,17 @@ public class ActiveConnectionRecord implements ConnectionRecord {
 
     @Override
     public GuacamoleSocket getActiveSocket() {
-        // STUB
-        return null;
+        return socket;
+    }
+
+    /**
+     * Associates the given GuacamoleSocket with this connection record.
+     *
+     * @param socket
+     *     The GuacamoleSocket to associate with this connection record.
+     */
+    public void setActiveSocket(GuacamoleSocket socket) {
+        this.socket = socket;
     }
 
 }