CP-54215: Limit to one VNC console session per VM/host

When the `limit_console_sessions` is true.
- Enforced a single active VNC console connection per VM/host
- Disable connection to websocket

Signed-off-by: Stephen Cheng <[email protected]>

Author: stephenchengCloud

ocaml/xapi/console.ml

@@ -34,6 +34,51 @@ type address =
 
 (* console is listening on a Unix domain socket *)
 
+(* This module enforces connection limits for VM consoles.
+   Depending on configuration, only one active connection per VM (including Dom0) is allowed,
+   or that connections are unlimited. *)
+module Connection_limit = struct
+  module VMSet = Set.Make (String)
+
+  let active_connections : VMSet.t ref = ref VMSet.empty
+
+  let mutex = Mutex.create ()
+
+  let add vm_id =
+    Mutex.lock mutex ;
+    active_connections := VMSet.add vm_id !active_connections ;
+    Mutex.unlock mutex
+
+  let remove vm_id =
+    Mutex.lock mutex ;
+    active_connections := VMSet.remove vm_id !active_connections ;
+    Mutex.unlock mutex
+
+  let exists vm_id =
+    Mutex.lock mutex ;
+    let present = VMSet.mem vm_id !active_connections in
+    Mutex.unlock mutex ; present
+
+  let can_add vm_id limit_console_sessions =
+    if not limit_console_sessions then
+      true
+    else
+      not (exists vm_id)
+end
+
+let connection_limit_reached __context vm_id =
+  let pool = Helpers.get_pool ~__context in
+  let limit_console_sessions =
+    Db.Pool.get_limit_console_sessions ~__context ~self:pool
+  in
+  if not (Connection_limit.can_add vm_id limit_console_sessions) then (
+    debug
+      "Console session limit reached: only one active session allowed for VM %s"
+      vm_id ;
+    true
+  ) else
+    false
+
 let string_of_address = function
   | Port x ->
       "localhost:" ^ string_of_int x
@@ -84,26 +129,43 @@ let address_of_console __context console : address option =
     ) ;
   address_option
 
-let real_proxy __context _ _ vnc_port s =
+let real_proxy __context console _ _ vnc_port s =
   try
-    Http_svr.headers s (Http.http_200_ok ()) ;
-    let vnc_sock =
-      match vnc_port with
-      | Port x ->
-          Unixext.open_connection_fd "127.0.0.1" x
-      | Path x ->
-          Unixext.open_connection_unix_fd x
-    in
-    (* Unixext.proxy closes fds itself so we must dup here *)
-    let s' = Unix.dup s in
-    debug "Connected; running proxy (between fds: %d and %d)"
-      (Unixext.int_of_file_descr vnc_sock)
-      (Unixext.int_of_file_descr s') ;
-    Unixext.proxy vnc_sock s' ;
-    debug "Proxy exited"
+    let vm = Db.Console.get_VM ~__context ~self:console in
+    let vm_id = Ref.string_of vm in
+    if connection_limit_reached __context vm_id then
+      Http_svr.headers s (Http.http_503_service_unavailable ())
+    else (
+      Http_svr.headers s (Http.http_200_ok ()) ;
+      let vnc_sock =
+        match vnc_port with
+        | Port x ->
+            Unixext.open_connection_fd "127.0.0.1" x
+        | Path x ->
+            Unixext.open_connection_unix_fd x
+      in
+      (* Unixext.proxy closes fds itself so we must dup here *)
+      let s' = Unix.dup s in
+      debug "Connected; running proxy (between fds: %d and %d)"
+        (Unixext.int_of_file_descr vnc_sock)
+        (Unixext.int_of_file_descr s') ;
+      Connection_limit.add vm_id ;
+      Unixext.proxy vnc_sock s' ;
+      Connection_limit.remove vm_id ;
+      debug "Proxy exited"
+    )
   with exn -> debug "error: %s" (ExnHelper.string_of_exn exn)
 
-let ws_proxy __context req protocol address s =
+let ws_proxy __context _ req protocol address s =
+  let pool = Helpers.get_pool ~__context in
+  let limit_console_sessions =
+    Db.Pool.get_limit_console_sessions ~__context ~self:pool
+  in
+  (* Disable connection via websocket if the limit is set *)
+  if limit_console_sessions = true then (
+    Http_svr.headers s (Http.http_503_service_unavailable ()) ;
+    ()
+  ) ;
   let addr = match address with Port p -> string_of_int p | Path p -> p in
   let protocol =
     match protocol with `rfb -> "rfb" | `vt100 -> "vt100" | `rdp -> "rdp"
@@ -247,7 +309,7 @@ let handler proxy_fn (req : Request.t) s _ =
       check_vm_is_running_here __context console ;
       match address_of_console __context console with
       | Some vnc_port ->
-          proxy_fn __context req protocol vnc_port s
+          proxy_fn __context console req protocol vnc_port s
       | None ->
           Http_svr.headers s (Http.http_404_missing ())
   )