console-proxy: close session if unable to connect to VNC server

Signed-off-by: Abhishek Kumar <[email protected]>

Author: shwstppr

services/console-proxy/server/src/main/java/com/cloud/consoleproxy/ConsoleProxyNoVncClient.java

@@ -109,7 +109,12 @@ public void run() {
                     String tunnelSession = param.getClientTunnelSession();
                     String websocketUrl = param.getWebsocketUrl();
 
-                    connectClientToVNCServer(tunnelUrl, tunnelSession, websocketUrl);
+                    if (!connectClientToVNCServer(tunnelUrl, tunnelSession, websocketUrl)) {
+                        logger.error("Failed to connect to VNC server, will close connection with client [{}] [IP: {}].", clientId, clientSourceIp);
+                        connectionAlive = false;
+                        session.close();
+                        return;
+                    }
                     authenticateToVNCServer(clientSourceIp);
 
                     // Track consecutive iterations with no data and sleep accordingly. Only used for NIO socket connections.
@@ -313,7 +318,7 @@ protected static byte[] rewriteServerNameInServerInit(byte[] serverInitBytes, St
      * - When websocketUrl is not empty -> connect to websocket
      * - Otherwise -> connect to TCP port on host directly
      */
-    private void connectClientToVNCServer(String tunnelUrl, String tunnelSession, String websocketUrl) {
+    private boolean connectClientToVNCServer(String tunnelUrl, String tunnelSession, String websocketUrl) {
         try {
             if (StringUtils.isNotBlank(websocketUrl)) {
                 logger.info(String.format("Connect to VNC over websocket URL: %s", websocketUrl));
@@ -337,7 +342,9 @@ private void connectClientToVNCServer(String tunnelUrl, String tunnelSession, St
             logger.info("Connection to VNC server has been established successfully.");
         } catch (Throwable e) {
             logger.error("Unexpected exception while connecting to VNC server.", e);
+            return false;
         }
+        return true;
     }
 
     private void setClientParam(ConsoleProxyClientParam param) {

services/console-proxy/server/src/main/java/com/cloud/consoleproxy/vnc/NoVncClient.java

@@ -88,15 +88,16 @@ public void connectTo(String host, int port, String path, String session, boolea
         setTunnelSocketStreams();
     }
 
-    public void connectTo(String host, int port) {
+    public void connectTo(String host, int port) throws IOException {
         // Connect to server
         logger.info("Connecting to VNC server {}:{} ...", host, port);
         try {
             NioSocket nioSocket = new NioSocket(host, port);
             this.nioSocketConnection = new NioSocketHandlerImpl(nioSocket);
-        } catch (Exception e) {
+        } catch (IOException e) {
             logger.error(String.format("Cannot create socket to host: %s and port %s: %s", host, port,
                     e.getMessage()), e);
+            throw e;
         }
     }
 

services/console-proxy/server/src/main/java/com/cloud/consoleproxy/vnc/network/NioSocket.java

@@ -36,7 +36,7 @@ public class NioSocket {
     private static final int CONNECTION_TIMEOUT_MILLIS = 3000;
     protected Logger logger = LogManager.getLogger(getClass());
 
-    private void initializeSocket() {
+    private void initializeSocket() throws IOException {
         try {
             socketChannel = SocketChannel.open();
             socketChannel.configureBlocking(false);
@@ -49,42 +49,40 @@ private void initializeSocket() {
             socketChannel.register(readSelector, SelectionKey.OP_READ);
         } catch (IOException e) {
             logger.error("Could not initialize NioSocket: " + e.getMessage(), e);
+            throw e;
         }
     }
 
-    private void waitForSocketSelectorConnected(Selector selector) {
-        try {
-            while (selector.select(CONNECTION_TIMEOUT_MILLIS) <= 0) {
-                logger.debug("Waiting for ready operations to connect to the socket");
-            }
-            Set<SelectionKey> keys = selector.selectedKeys();
-            for (SelectionKey selectionKey: keys) {
-                if (selectionKey.isConnectable()) {
-                    if (socketChannel.isConnectionPending()) {
-                        socketChannel.finishConnect();
-                    }
-                    logger.debug("Connected to the socket");
-                    break;
+    private void waitForSocketSelectorConnected(Selector selector) throws IOException {
+        while (selector.select(CONNECTION_TIMEOUT_MILLIS) <= 0) {
+            logger.debug("Waiting for ready operations to connect to the socket");
+        }
+        Set<SelectionKey> keys = selector.selectedKeys();
+        for (SelectionKey selectionKey: keys) {
+            if (selectionKey.isConnectable()) {
+                if (socketChannel.isConnectionPending()) {
+                    socketChannel.finishConnect();
                 }
+                logger.debug("Connected to the socket");
+                break;
             }
-        } catch (IOException e) {
-            logger.error(String.format("Error waiting for socket selector ready: %s", e.getMessage()), e);
         }
     }
 
-    private void connectSocket(String host, int port) {
+    private void connectSocket(String host, int port) throws IOException {
         try {
             socketChannel.connect(new InetSocketAddress(host, port));
             Selector selector = Selector.open();
             socketChannel.register(selector, SelectionKey.OP_CONNECT);
 
             waitForSocketSelectorConnected(selector);
         } catch (IOException e) {
-            logger.error(String.format("Error creating NioSocket to %s:%s: %s", host, port, e.getMessage()), e);
+            logger.error("Error connecting NioSocket to {}:{}: {}", host, port, e.getMessage(), e);
+            throw e;
         }
     }
 
-    public NioSocket(String host, int port) {
+    public NioSocket(String host, int port) throws IOException {
         initializeSocket();
         connectSocket(host, port);
     }