console-proxy: stop busy-spinning in NoVncClient.waitForNoVNCReply

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

@@ -178,8 +178,14 @@ public void run() {
                     }
                     logger.info("Connection with client [{}] [IP: {}] is dead.", clientId, clientSourceIp);
                     ConsoleProxy.grantReconnectWindowForSessionAndClientIp(sessionUuid, clientSourceIp);
-                } catch (IOException e) {
-                    logger.error("Error on VNC client", e);
+                } catch (IOException | RuntimeException e) {
+                    if (connectionAlive) {
+                        logger.error("Error on VNC client [{}] [IP: {}], closing the connection.", clientId, clientSourceIp, e);
+                    } else {
+                        logger.debug("VNC client [{}] [IP: {}] stopped after its connection was closed: {}", clientId, clientSourceIp, e.getMessage());
+                    }
+                    closeClient();
+                    session.close();
                 }
             }
 
@@ -253,8 +259,9 @@ protected void handshakeSecurityResult(int secType) {
 
         client.processSecurityResultMsg();
         byte[] securityResultToClient = new byte[] { 0, 0, 0, 0 };
-        sendMessageToVNCClient(securityResultToClient, 4);
+        // Set the flag before sending, as noVNC may reply before sendMessageToVNCClient() returns
         client.setWaitForNoVnc(true);
+        sendMessageToVNCClient(securityResultToClient, 4);
     }
 
     protected int handshakeSecurity() {
@@ -288,8 +295,8 @@ private void initialisationPhase() {
                 client.isTLSConnectionEstablished() ? "(TLS backend)" : "");
         byte[] bytesServerInit = rewriteServerNameInServerInit(serverInitByteArray, displayNameForVM);
 
-        sendMessageToVNCClient(bytesServerInit, bytesServerInit.length);
         client.setWaitForNoVnc(true);
+        sendMessageToVNCClient(bytesServerInit, bytesServerInit.length);
         client.waitForNoVNCReply();
     }
 

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

@@ -30,6 +30,7 @@
 import java.security.spec.KeySpec;
 import java.util.Arrays;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 
 import com.cloud.consoleproxy.util.RawHTTP;
 import com.cloud.consoleproxy.vnc.network.NioSocket;
@@ -59,6 +60,8 @@
 public class NoVncClient {
     protected Logger logger = LogManager.getLogger(getClass());
 
+    protected static final long NOVNC_REPLY_TIMEOUT_MS = 30000L;
+
     private Socket socket;
     private DataInputStream is;
     private DataOutputStream os;
@@ -130,6 +133,10 @@ public void proxyMsgOverWebSocketConnection(ByteBuffer msg) {
     }
 
     public void close() {
+        synchronized (lock) {
+            closed = true;
+            lock.notifyAll();
+        }
         if (nioSocketConnection != null) {
             nioSocketConnection.close();
         }
@@ -413,13 +420,32 @@ public ByteBuffer handshakeProtocolVersion() {
         return verStr;
     }
 
+    /**
+     * Blocks until the noVNC client replies (see {@link #setWaitForNoVnc(boolean)}). Gives up with a
+     * {@link CloudRuntimeException} if the client is closed or no reply arrives within {@link #NOVNC_REPLY_TIMEOUT_MS}.
+     */
     public void waitForNoVNCReply() {
-        int cycles = 0;
-        while (isWaitForNoVnc()) {
-            cycles++;
-        }
-        if (logger.isDebugEnabled()) {
-            logger.debug(String.format("Waited %d cycles for NoVnc", cycles));
+        waitForNoVNCReply(NOVNC_REPLY_TIMEOUT_MS);
+    }
+
+    protected void waitForNoVNCReply(long timeoutMs) {
+        long deadline = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(timeoutMs);
+        synchronized (lock) {
+            while (waitForNoVnc) {
+                if (closed) {
+                    throw new CloudRuntimeException("Connection closed while waiting for the noVNC reply");
+                }
+                long remainingMs = TimeUnit.NANOSECONDS.toMillis(deadline - System.nanoTime());
+                if (remainingMs <= 0) {
+                    throw new CloudRuntimeException(String.format("Timed out after %d ms waiting for the noVNC reply", timeoutMs));
+                }
+                try {
+                    lock.wait(remainingMs);
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    throw new CloudRuntimeException("Interrupted while waiting for the noVNC reply", e);
+                }
+            }
         }
     }
 
@@ -468,6 +494,7 @@ public int handshakeSecurityType() {
     public void setWaitForNoVnc(boolean val) {
         synchronized (lock) {
             this.waitForNoVnc = val;
+            lock.notifyAll();
         }
     }
 
@@ -478,6 +505,7 @@ public boolean isWaitForNoVnc() {
     }
 
     private boolean waitForNoVnc = false;
+    private boolean closed = false;
 
     private Pair<Boolean, String> processSecurityResultType(int authResult) {
         boolean result = false;

services/console-proxy/server/src/test/java/com/cloud/consoleproxy/ConsoleProxyNoVncClientTest.java

@@ -16,8 +16,25 @@
 // under the License.
 package com.cloud.consoleproxy;
 
+import java.io.OutputStream;
+import java.net.InetAddress;
+import java.net.ServerSocket;
+import java.net.Socket;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+import org.eclipse.jetty.websocket.api.RemoteEndpoint;
+import org.eclipse.jetty.websocket.api.Session;
 import org.junit.Assert;
 import org.junit.Test;
+import org.mockito.InOrder;
+import org.mockito.Mockito;
+import org.springframework.test.util.ReflectionTestUtils;
+
+import com.cloud.consoleproxy.vnc.NoVncClient;
+import com.cloud.consoleproxy.vnc.RfbConstants;
 
 public class ConsoleProxyNoVncClientTest {
     @Test
@@ -29,4 +46,87 @@ public void rewriteServerNameInServerInitTest() {
         byte[] expectedBytes = new byte[]{4, 0, 3, 0, 32, 24, 0, 1, 0, -1, 0, -1, 0, -1, 16, 8, 0, 0, 0, 0, 0, 0, 0, 22, 115, 101, 114, 118, 101, 114, 49, 50, 51, 44, 32, 98, 97, 99, 107, 101, 110, 100, 58, 84, 76, 83};
         Assert.assertArrayEquals(newServerInit, expectedBytes);
     }
+
+    @Test
+    public void authenticateVNCServerThroughNioSocketExpectsNoVncReplyBeforeSendingTheMessage() throws Exception {
+        Session session = Mockito.mock(Session.class);
+        RemoteEndpoint remote = Mockito.mock(RemoteEndpoint.class);
+        Mockito.when(session.getRemote()).thenReturn(remote);
+        NoVncClient client = Mockito.mock(NoVncClient.class);
+        Mockito.when(client.handshakeProtocolVersion()).thenReturn(ByteBuffer.wrap("RFB 003.008\n".getBytes(StandardCharsets.US_ASCII)));
+        Mockito.when(client.handshakeSecurityType()).thenReturn(RfbConstants.NO_AUTH);
+        Mockito.when(client.readServerInit()).thenReturn(new byte[24]);
+        ConsoleProxyClientParam param = new ConsoleProxyClientParam();
+        param.setClientDisplayName("vm");
+
+        ConsoleProxyNoVncClient viewer = new ConsoleProxyNoVncClient(session);
+        ReflectionTestUtils.setField(viewer, "client", client);
+        ReflectionTestUtils.setField(viewer, "clientParam", param);
+        viewer.authenticateVNCServerThroughNioSocket();
+
+        // A reply flagged after the message is sent can arrive first and be lost
+        InOrder inOrder = Mockito.inOrder(client, remote);
+        inOrder.verify(client).processSecurityResultMsg();
+        inOrder.verify(client).setWaitForNoVnc(true);
+        inOrder.verify(remote).sendBytes(Mockito.argThat(msg -> msg.remaining() == 4));
+        inOrder.verify(client).waitForNoVNCReply();
+        inOrder.verify(client).readServerInit();
+        inOrder.verify(client).setWaitForNoVnc(true);
+        inOrder.verify(remote).sendBytes(Mockito.argThat(msg -> msg.remaining() > 24));
+        inOrder.verify(client).waitForNoVNCReply();
+    }
+
+    @Test(timeout = 20000)
+    public void handshakeFailureClosesTheSession() throws Exception {
+        try (ServerSocket vncServer = new ServerSocket(0, 1, InetAddress.getLoopbackAddress())) {
+            Session session = Mockito.mock(Session.class);
+            ConsoleProxyNoVncClient viewer = new ConsoleProxyNoVncClient(session);
+            viewer.initClient(createClientParam(vncServer));
+
+            // The VNC server goes away before sending its protocol version
+            vncServer.accept().close();
+
+            Mockito.verify(session, Mockito.timeout(10000)).close();
+            Assert.assertFalse(viewer.isHostConnected());
+        }
+    }
+
+    @Test(timeout = 20000)
+    public void closingTheViewerWhileWaitingForNoVncStopsTheWorker() throws Exception {
+        try (ServerSocket vncServer = new ServerSocket(0, 1, InetAddress.getLoopbackAddress())) {
+            Session session = Mockito.mock(Session.class);
+            RemoteEndpoint remote = Mockito.mock(RemoteEndpoint.class);
+            Mockito.when(session.getRemote()).thenReturn(remote);
+            CountDownLatch versionSentToNoVnc = new CountDownLatch(1);
+            Mockito.doAnswer(invocation -> {
+                versionSentToNoVnc.countDown();
+                return null;
+            }).when(remote).sendBytes(Mockito.any(ByteBuffer.class));
+
+            ConsoleProxyNoVncClient viewer = new ConsoleProxyNoVncClient(session);
+            viewer.initClient(createClientParam(vncServer));
+
+            try (Socket vncConnection = vncServer.accept()) {
+                OutputStream out = vncConnection.getOutputStream();
+                out.write("RFB 003.008\n".getBytes(StandardCharsets.US_ASCII));
+                out.flush();
+                Assert.assertTrue(versionSentToNoVnc.await(10, TimeUnit.SECONDS));
+
+                // noVNC never replies, then its websocket is closed (ConsoleProxyNoVNCHandler.onClose)
+                viewer.closeClient();
+
+                Mockito.verify(session, Mockito.timeout(10000)).close();
+                vncConnection.setSoTimeout(10000);
+                Assert.assertEquals(-1, vncConnection.getInputStream().read());
+            }
+        }
+    }
+
+    private static ConsoleProxyClientParam createClientParam(ServerSocket vncServer) {
+        ConsoleProxyClientParam param = new ConsoleProxyClientParam();
+        param.setClientHostAddress(vncServer.getInetAddress().getHostAddress());
+        param.setClientHostPort(vncServer.getLocalPort());
+        param.setClientIp("127.0.0.1");
+        return param;
+    }
 }

services/console-proxy/server/src/test/java/com/cloud/consoleproxy/vnc/NoVncClientTest.java

@@ -0,0 +1,94 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+package com.cloud.consoleproxy.vnc;
+
+import com.cloud.utils.exception.CloudRuntimeException;
+import org.junit.Assert;
+import org.junit.Test;
+
+public class NoVncClientTest {
+
+    @Test(timeout = 10000)
+    public void waitForNoVNCReplyTimesOutWhenNoReplyArrives() {
+        NoVncClient client = new NoVncClient();
+        client.setWaitForNoVnc(true);
+        long start = System.nanoTime();
+        Assert.assertThrows(CloudRuntimeException.class, () -> client.waitForNoVNCReply(200));
+        long elapsedMs = (System.nanoTime() - start) / 1000000;
+        Assert.assertTrue("Gave up after " + elapsedMs + " ms", elapsedMs >= 150);
+    }
+
+    @Test(timeout = 10000, expected = CloudRuntimeException.class)
+    public void waitForNoVNCReplyFailsFastWhenAlreadyClosed() {
+        NoVncClient client = new NoVncClient();
+        client.close();
+        client.setWaitForNoVnc(true);
+        client.waitForNoVNCReply();
+    }
+
+    @Test
+    public void waitForNoVNCReplyReturnsImmediatelyWhenNotWaiting() {
+        NoVncClient client = new NoVncClient();
+        client.waitForNoVNCReply();
+    }
+
+    @Test(timeout = 10000)
+    public void waitForNoVNCReplyReturnsWhenReplyArrives() throws Exception {
+        NoVncClient client = new NoVncClient();
+        client.setWaitForNoVnc(true);
+        Thread replier = new Thread(() -> {
+            try {
+                Thread.sleep(100);
+            } catch (InterruptedException ignored) {
+                Thread.currentThread().interrupt();
+            }
+            client.setWaitForNoVnc(false);
+        });
+        replier.start();
+        client.waitForNoVNCReply();
+        replier.join();
+        Assert.assertFalse(client.isWaitForNoVnc());
+    }
+
+    @Test(timeout = 10000, expected = CloudRuntimeException.class)
+    public void waitForNoVNCReplyFailsWhenClientIsClosed() {
+        NoVncClient client = new NoVncClient();
+        client.setWaitForNoVnc(true);
+        Thread closer = new Thread(() -> {
+            try {
+                Thread.sleep(100);
+            } catch (InterruptedException ignored) {
+                Thread.currentThread().interrupt();
+            }
+            client.close();
+        });
+        closer.start();
+        client.waitForNoVNCReply();
+    }
+
+    @Test(timeout = 10000, expected = CloudRuntimeException.class)
+    public void waitForNoVNCReplyFailsWhenInterrupted() {
+        NoVncClient client = new NoVncClient();
+        client.setWaitForNoVnc(true);
+        Thread.currentThread().interrupt();
+        try {
+            client.waitForNoVNCReply();
+        } finally {
+            Thread.interrupted();
+        }
+    }
+}