From 075e880acc7aa115f0a4b27fed3f182f8108affa Mon Sep 17 00:00:00 2001
From: Nick Couchman <vnick@apache.org>
Date: Wed, 31 May 2017 15:22:31 -0400
Subject: [PATCH] GUACAMOLE-102: Deal with weights of 0, and properly dispose
 of connections with negative weights.

---
 .../jdbc/connection/ModeledConnection.java    |  3 +-
 .../RestrictedGuacamoleTunnelService.java     | 41 ++++++++++---------
 2 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/extensions/guacamole-auth-jdbc/modules/guacamole-auth-jdbc-base/src/main/java/org/apache/guacamole/auth/jdbc/connection/ModeledConnection.java b/extensions/guacamole-auth-jdbc/modules/guacamole-auth-jdbc-base/src/main/java/org/apache/guacamole/auth/jdbc/connection/ModeledConnection.java
index 1044a6202..414a8a43f 100644
--- a/extensions/guacamole-auth-jdbc/modules/guacamole-auth-jdbc-base/src/main/java/org/apache/guacamole/auth/jdbc/connection/ModeledConnection.java
+++ b/extensions/guacamole-auth-jdbc/modules/guacamole-auth-jdbc-base/src/main/java/org/apache/guacamole/auth/jdbc/connection/ModeledConnection.java
@@ -409,6 +409,7 @@ public class ModeledConnection extends ModeledChildDirectoryObject<ConnectionMod
             port             != null ? port             : defaultConfig.getPort(),
             encryptionMethod != null ? encryptionMethod : defaultConfig.getEncryptionMethod()
         );
+    }
 
     /** 
      * Returns the weight of the connection, or the default.
@@ -417,7 +418,7 @@ public class ModeledConnection extends ModeledChildDirectoryObject<ConnectionMod
      *     The weight of the connection.
      *  
      */
-    public int getConnectionWeight() {
+    public Integer getConnectionWeight() {
 
         // Return the connection weight
         return getModel().getConnectionWeight();
diff --git a/extensions/guacamole-auth-jdbc/modules/guacamole-auth-jdbc-base/src/main/java/org/apache/guacamole/auth/jdbc/tunnel/RestrictedGuacamoleTunnelService.java b/extensions/guacamole-auth-jdbc/modules/guacamole-auth-jdbc-base/src/main/java/org/apache/guacamole/auth/jdbc/tunnel/RestrictedGuacamoleTunnelService.java
index b41ec91bc..49af081d0 100644
--- a/extensions/guacamole-auth-jdbc/modules/guacamole-auth-jdbc-base/src/main/java/org/apache/guacamole/auth/jdbc/tunnel/RestrictedGuacamoleTunnelService.java
+++ b/extensions/guacamole-auth-jdbc/modules/guacamole-auth-jdbc-base/src/main/java/org/apache/guacamole/auth/jdbc/tunnel/RestrictedGuacamoleTunnelService.java
@@ -23,6 +23,7 @@ import com.google.common.collect.ConcurrentHashMultiset;
 import com.google.inject.Inject;
 import com.google.inject.Singleton;
 import java.util.Arrays;
+import java.util.Iterator;
 import java.util.Comparator;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -46,11 +47,6 @@ import org.slf4j.LoggerFactory;
 public class RestrictedGuacamoleTunnelService
     extends AbstractGuacamoleTunnelService {
 
-    /**
-     * Logger for this class.
-     */
-    private static final Logger logger = LoggerFactory.getLogger(RestrictedGuacamoleTunnelService.class);
-
     /**
      * The environment of the Guacamole server.
      */
@@ -180,6 +176,14 @@ public class RestrictedGuacamoleTunnelService
         // Get username
         String username = user.getIdentifier();
 
+        // Remove connections where weight < 0
+        Iterator<ModeledConnection> i = connections.iterator();
+        while(i.hasNext()) {
+            Integer weight = i.next().getConnectionWeight();
+            if (weight != null && weight.intValue() < 0)
+                i.remove();
+        }
+
         // Sort connections in ascending order of usage
         ModeledConnection[] sortedConnections = connections.toArray(new ModeledConnection[connections.size()]);
         Arrays.sort(sortedConnections, new Comparator<ModeledConnection>() {
@@ -187,22 +191,21 @@ public class RestrictedGuacamoleTunnelService
             @Override
             public int compare(ModeledConnection a, ModeledConnection b) {
                 
-                logger.debug("Calculating weights for connections {} and {}.", a.getName(), b.getName());
-                int cw = 0;
-                int weightA = a.getConnectionWeight();
-                // If the weight is null, we go ahead and sort, anyway
-                if (weightA == null)
-                    weightA = 0;
+                int weightA, weightB;
+                // Check if weight of a is null, assign 1 if it is.
+                if (a.getConnectionWeight() == null)
+                    weightA = 1;
+                else
+                    weightA = a.getConnectionWeight().intValue() + 1;
 
-                // If the weight is null, we go ahead and sort, anyway
-                int weightB = b.getConnectionWeight();
-                if (weightB == null)
-                    weightB = 0;
+                // Check if weight of b is null, assign 1 if it is.
+                if (b.getConnectionWeight() == null)
+                    weightB = 1;
+                else
+                    weightB = b.getConnectionWeight().intValue() + 1;
 
-                int connsA = getActiveConnections(a).size();
-                int connsB = getActiveConnections(b).size();
-                logger.debug("Connection {} has computed weight of {}.", a.getName(), connsA * 10000 / weightA);
-                logger.debug("Connection {} has computed weight of {}.", b.getName(), connsB * 10000 / weightB);
+                int connsA = getActiveConnections(a).size() + 1;
+                int connsB = getActiveConnections(b).size() + 1;
 
                 return (connsA * 10000 / weightA) - (connsB * 10000 / weightB);