Solves a potential panic (#14323)

Splits `close()` into `close()` and `cleanup()`. `cleanup()` is only ever called when both the input and output streaming goroutines have returned, so panics due to using freed memory should no longer be possible

lib/srv/desktop/rdp/rdpclient/client.go

@@ -150,7 +150,6 @@ func New(ctx context.Context, cfg Config) (*Client, error) {
 		cfg:           cfg,
 		readyForInput: 0,
 	}
-	c.handle = cgo.NewHandle(c)
 
 	if err := c.readClientUsername(); err != nil {
 		return nil, trace.Wrap(err)
@@ -167,14 +166,23 @@ func New(ctx context.Context, cfg Config) (*Client, error) {
 // Run starts the rdp client and blocks until the client disconnects,
 // then ensures the cleanup is run.
 func (c *Client) Run(ctx context.Context) error {
-	defer c.close()
+	defer c.cleanup()
+
+	c.handle = cgo.NewHandle(c)
 
 	if err := c.connect(ctx); err != nil {
 		return trace.Wrap(err)
 	}
 	c.start()
+
+	// Hang until input and output streaming
+	// goroutines both finish.
 	c.wg.Wait()
 
+	// Both goroutines have finished, it's now
+	// safe for the deferred c.cleanup() call to
+	// clean up the memory.
+
 	return nil
 }
 
@@ -246,6 +254,7 @@ func (c *Client) connect(ctx context.Context) error {
 		return trace.ConnectionProblem(nil, "RDP connection failed")
 	}
 	c.rustClient = res.client
+
 	return nil
 }
 
@@ -259,16 +268,13 @@ func (c *Client) start() {
 		defer c.close()
 		defer c.cfg.Log.Info("RDP output streaming finished")
 
+		c.cfg.Log.Info("RDP output streaming starting")
+
 		// C.read_rdp_output blocks for the duration of the RDP connection and
 		// calls handle_bitmap repeatedly with the incoming bitmaps.
 		if errCode := C.read_rdp_output(c.rustClient); errCode != C.ErrCodeSuccess {
 			c.cfg.Log.Warningf("Failed reading RDP output frame: %v", errCode)
-
-			// close the TDP connection to the browser
-			// (without this the input streaming goroutine will hang
-			// waiting for user input)
 			c.cfg.Conn.SendError("There was an error reading data from the Windows Desktop")
-			c.cfg.Conn.Close()
 		}
 	}()
 
@@ -279,6 +285,8 @@ func (c *Client) start() {
 		defer c.close()
 		defer c.cfg.Log.Info("TDP input streaming finished")
 
+		c.cfg.Log.Info("TDP input streaming starting")
+
 		// Remember mouse coordinates to send them with all CGOPointer events.
 		var mouseX, mouseY uint32
 		for {
@@ -519,29 +527,43 @@ func (c *Client) sharedDirectoryInfoRequest(req tdp.SharedDirectoryInfoRequest)
 	return C.ErrCodeSuccess
 }
 
-// close frees the memory of the cgo.Handle,
-// closes the RDP client connection,
-// and frees the Rust client.
+// close closes the RDP client connection and
+// the TDP connection to the browser.
 func (c *Client) close() {
 	c.closeOnce.Do(func() {
 		// Ensure the RDP connection is closed
 		if errCode := C.close_rdp(c.rustClient); errCode != C.ErrCodeSuccess {
 			c.cfg.Log.Warningf("error closing the RDP connection")
+		} else {
+			c.cfg.Log.Debug("RDP connection closed successfully")
 		}
 
-		// Let the Rust side free its data
-		C.free_rdp(c.rustClient)
-
-		// Release the memory of the cgo.Handle
-		c.handle.Delete()
-
 		// Ensure the TDP connection is closed
 		if err := c.cfg.Conn.Close(); err != nil {
 			c.cfg.Log.Warningf("error closing the TDP connection: %v", err)
+		} else {
+			c.cfg.Log.Debug("TDP connection closed successfully")
 		}
 	})
 }
 
+// cleanup frees the Rust client and
+// frees the memory of the cgo.Handle.
+// This function should only be called
+// once per Client.
+func (c *Client) cleanup() {
+	// Let the Rust side free its data
+	if c.rustClient != nil {
+		C.free_rdp(c.rustClient)
+	}
+
+	// Release the memory of the cgo.Handle
+	if c.handle != 0 {
+		c.handle.Delete()
+	}
+
+}
+
 // GetClientLastActive returns the time of the last recorded activity.
 // For RDP, "activity" is defined as user-input messages
 // (mouse move, button press, etc.)

lib/srv/desktop/rdp/rdpclient/librdprs.h

@@ -57,7 +57,7 @@ typedef enum CGOPointerWheel {
  * - connect_rdp creates it on the heap, grabs a raw pointer and returns in to Go
  * - most other exported rdp functions take the raw pointer, convert it to a reference for use
  *   without dropping the Client
- * - close_rdp takes the raw pointer and drops it
+ * - free_rdp takes the raw pointer and drops it
  *
  * All of the exported rdp functions could run concurrently, so the rdp_client is synchronized.
  * tcp_fd is only set in connect_rdp and used as read-only afterwards, so it does not need

lib/srv/desktop/rdp/rdpclient/src/lib.rs

@@ -57,7 +57,7 @@ pub extern "C" fn init() {
 /// - connect_rdp creates it on the heap, grabs a raw pointer and returns in to Go
 /// - most other exported rdp functions take the raw pointer, convert it to a reference for use
 ///   without dropping the Client
-/// - close_rdp takes the raw pointer and drops it
+/// - free_rdp takes the raw pointer and drops it
 ///
 /// All of the exported rdp functions could run concurrently, so the rdp_client is synchronized.
 /// tcp_fd is only set in connect_rdp and used as read-only afterwards, so it does not need
@@ -800,11 +800,9 @@ pub unsafe extern "C" fn close_rdp(client_ptr: *mut Client) -> CGOErrCode {
             return cgo_error;
         }
     };
-    if let Err(e) = client.rdp_client.lock().unwrap().shutdown() {
-        error!("failed writing RDP keyboard event: {:?}", e);
-        CGOErrCode::ErrCodeFailure
-    } else {
-        CGOErrCode::ErrCodeSuccess
+    match client.rdp_client.lock().unwrap().shutdown() {
+        Err(_) => CGOErrCode::ErrCodeFailure,
+        Ok(_) => CGOErrCode::ErrCodeSuccess,
     }
 }