Add latency detector for desktop sessions

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

@@ -74,8 +74,10 @@ import (
 	"encoding/binary"
 	"fmt"
 	"log/slog"
+	"net"
 	"os"
 	"runtime/cgo"
+	"strconv"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -414,6 +416,9 @@ func (c *Client) startInputStreaming(stopCh chan struct{}) error {
 	c.cfg.Logger.InfoContext(context.Background(), "TDP input streaming starting")
 	defer c.cfg.Logger.InfoContext(context.Background(), "TDP input streaming finished")
 
+	// we will disable ping only if the env var is truthy
+	disableDesktopPing, _ := strconv.ParseBool(os.Getenv("TELEPORT_DISABLE_DESKTOP_LATENCY_DETECTOR_PING"))
+
 	var withheldResize *tdp.ClientScreenSpec
 	for {
 		select {
@@ -432,6 +437,22 @@ func (c *Client) startInputStreaming(stopCh chan struct{}) error {
 			c.cfg.Logger.WarnContext(context.Background(), "Failed reading TDP input message", "error", err)
 			return err
 		}
+		if m, ok := msg.(tdp.Ping); ok {
+			// Upon receiving a ping message, we make a connection
+			// to the host and send the same message back to the proxy.
+			// The proxy will then compute the round trip time.
+			if !disableDesktopPing {
+				conn, err := net.Dial("tcp", c.cfg.Addr)
+				if err == nil {
+					conn.Close()
+				}
+			}
+			if err := c.cfg.Conn.WriteMessage(m); err != nil {
+				c.cfg.Logger.WarnContext(context.Background(), "Failed writing TDP ping message", "error", err)
+				return err
+			}
+			continue
+		}
 
 		if atomic.LoadUint32(&c.readyForInput) == 0 {
 			switch m := msg.(type) {

lib/srv/desktop/tdp/proto.go

@@ -33,6 +33,7 @@ import (
 	"image/png"
 	"io"
 
+	"github.com/google/uuid"
 	"github.com/gravitational/trace"
 
 	authproto "github.com/gravitational/teleport/api/client/proto"
@@ -82,6 +83,8 @@ const (
 	TypeSyncKeys                        = MessageType(32)
 	TypeSharedDirectoryTruncateRequest  = MessageType(33)
 	TypeSharedDirectoryTruncateResponse = MessageType(34)
+	TypeLatencyStats                    = MessageType(35)
+	TypePing                            = MessageType(36)
 )
 
 // Message is a Go representation of a desktop protocol message.
@@ -182,6 +185,8 @@ func decodeMessage(firstByte byte, in byteReader) (Message, error) {
 		return decodeSharedDirectoryTruncateRequest(in)
 	case TypeSharedDirectoryTruncateResponse:
 		return decodeSharedDirectoryTruncateResponse(in)
+	case TypePing:
+		return decodePing(in)
 	default:
 		return nil, trace.BadParameter("unsupported desktop protocol message type %d", firstByte)
 	}
@@ -1631,6 +1636,44 @@ func decodeSharedDirectoryTruncateResponse(in io.Reader) (SharedDirectoryTruncat
 	return res, err
 }
 
+// LatencyStats is used to report the latency of the connection(s) to the client.
+type LatencyStats struct {
+	ClientLatency uint32
+	ServerLatency uint32
+}
+
+func (l LatencyStats) Encode() ([]byte, error) {
+	buf := new(bytes.Buffer)
+	buf.WriteByte(byte(TypeLatencyStats))
+	writeUint32(buf, l.ClientLatency)
+	writeUint32(buf, l.ServerLatency)
+	return buf.Bytes(), nil
+}
+
+// Ping is used to measure the latency of the connection(s) between proxy and desktop (includes
+// latency between proxy and Windows Desktop Service and between WDS and desktop).
+type Ping struct {
+
+	// UUID is used to correlate message send by proxy and received from the Windows Desktop Service
+	UUID uuid.UUID
+}
+
+func (p Ping) Encode() ([]byte, error) {
+	buf := new(bytes.Buffer)
+	buf.WriteByte(byte(TypePing))
+	buf.Write(p.UUID[:])
+	return buf.Bytes(), nil
+}
+
+func decodePing(in io.Reader) (Ping, error) {
+	var ping Ping
+	_, err := io.ReadFull(in, ping.UUID[:])
+	if err != nil {
+		return ping, trace.Wrap(err)
+	}
+	return ping, nil
+}
+
 // encodeString encodes strings for TDP. Strings are encoded as UTF-8 with
 // a 32-bit length prefix (in bytes):
 // https://github.com/gravitational/teleport/blob/master/rfd/0037-desktop-access-protocol.md#field-types

lib/web/desktop.go

@@ -30,9 +30,12 @@ import (
 	"net/http"
 	"sync"
 
+	"github.com/google/uuid"
 	"github.com/gorilla/websocket"
 	"github.com/gravitational/trace"
+	"github.com/jonboulle/clockwork"
 	"github.com/julienschmidt/httprouter"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/gravitational/teleport/api/client/proto"
 	"github.com/gravitational/teleport/api/constants"
@@ -47,6 +50,7 @@ import (
 	"github.com/gravitational/teleport/lib/reversetunnelclient"
 	"github.com/gravitational/teleport/lib/srv/desktop/tdp"
 	"github.com/gravitational/teleport/lib/utils"
+	"github.com/gravitational/teleport/lib/utils/diagnostics/latency"
 	logutils "github.com/gravitational/teleport/lib/utils/log"
 )
 
@@ -194,7 +198,7 @@ func (h *Handler) createDesktopConnection(
 		clientSrcAddr: clientSrcAddr,
 		clientDstAddr: clientDstAddr,
 	}
-	serviceConn, _, err := c.connectToWindowsService(ctx, clusterName, validServiceIDs)
+	serviceConn, version, err := c.connectToWindowsService(ctx, clusterName, validServiceIDs)
 	if err != nil {
 		return sendTDPError(trace.Wrap(err, "cannot connect to Windows Desktop Service"))
 	}
@@ -233,7 +237,7 @@ func (h *Handler) createDesktopConnection(
 	// proxyWebsocketConn hangs here until connection is closed
 	handleProxyWebsocketConnErr(
 		ctx,
-		proxyWebsocketConn(ws, serviceConnTLS),
+		proxyWebsocketConn(ctx, ws, serviceConnTLS, log, version),
 		log,
 	)
 
@@ -535,19 +539,108 @@ func (c *connector) tryConnect(ctx context.Context, clusterName, desktopServiceI
 	return conn, ver, trace.Wrap(err)
 }
 
+// desktopPinger measures latency between proxy and the desktop by sending tdp.Ping messages
+// Windows Desktop Service and measuring the time it takes to receive message with the same UUID back.
+type desktopPinger struct {
+	wds net.Conn
+	ch  <-chan tdp.Ping
+}
+
+func (d desktopPinger) Ping(ctx context.Context) error {
+	ping := tdp.Ping{
+		UUID: uuid.New(),
+	}
+	buf, err := ping.Encode()
+	if err != nil {
+		return trace.Wrap(err)
+	}
+	_, err = d.wds.Write(buf)
+	if err != nil {
+		return trace.Wrap(err)
+	}
+	for {
+		select {
+		case pong := <-d.ch:
+			if pong.UUID == ping.UUID {
+				return nil
+			}
+		case <-ctx.Done():
+			return trace.Wrap(ctx.Err())
+		}
+	}
+}
+
 // proxyWebsocketConn does a bidrectional copy between the websocket
 // connection to the browser (ws) and the mTLS connection to Windows
 // Desktop Serivce (wds)
-func proxyWebsocketConn(ws *websocket.Conn, wds net.Conn) error {
+func proxyWebsocketConn(ctx context.Context, ws *websocket.Conn, wds net.Conn, log *slog.Logger, version string) error {
+	ctx, cancel := context.WithCancel(ctx)
 	var closeOnce sync.Once
 	close := func() {
+		cancel()
 		ws.Close()
 		wds.Close()
 	}
 
-	errs := make(chan error, 2)
+	tdpMessagesToSend := make(chan tdp.Message)
+
+	latencySupported, err := utils.MinVerWithoutPreRelease(version, "17.5.0")
+	if err != nil {
+		return trace.Wrap(err)
+	}
+
+	pings := make(chan tdp.Ping)
+
+	if latencySupported {
+		pinger := desktopPinger{
+			wds: wds,
+			ch:  pings,
+		}
+
+		go monitorLatency(ctx, clockwork.NewRealClock(), ws, pinger,
+			latency.ReporterFunc(func(ctx context.Context, stats latency.Statistics) error {
+				tdpMessagesToSend <- tdp.LatencyStats{
+					ClientLatency: uint32(stats.Client),
+					ServerLatency: uint32(stats.Server),
+				}
+				return nil
+			}),
+		)
+
+	}
+
+	var errs errgroup.Group
+
+	// run a goroutine to pick TDP messages up from a channel and send
+	// them to the browser
+	errs.Go(func() error {
+		for msg := range tdpMessagesToSend {
+			if ping, ok := msg.(tdp.Ping); ok {
+				pings <- ping
+				continue
+			}
+			if ls, ok := msg.(tdp.LatencyStats); ok {
+				log.DebugContext(ctx, "sending latency stats", "client", ls.ClientLatency, "server", ls.ServerLatency)
+			}
+			encoded, err := msg.Encode()
+			if err != nil {
+				return err
+			}
+
+			err = ws.WriteMessage(websocket.BinaryMessage, encoded)
+			if utils.IsOKNetworkError(err) {
+				return err
+			}
+			if err != nil {
+				return err
+			}
+		}
+		return nil
+	})
 
-	go func() {
+	// run a second goroutine to read TDP messages from the Windows
+	// agent and write them to our send channel
+	errs.Go(func() error {
 		defer closeOnce.Do(close)
 
 		// we avoid using io.Copy here, as we want to make sure
@@ -563,8 +656,7 @@ func proxyWebsocketConn(ws *websocket.Conn, wds net.Conn) error {
 		for {
 			msg, err := tc.ReadMessage()
 			if utils.IsOKNetworkError(err) {
-				errs <- nil
-				return
+				return err
 			} else if err != nil {
 				isFatal := tdp.IsFatalErr(err)
 				severity := tdp.SeverityError
@@ -585,58 +677,38 @@ func proxyWebsocketConn(ws *websocket.Conn, wds net.Conn) error {
 				if sendErr != nil {
 					err = sendErr
 				}
-				errs <- err
-				return
-			}
-			encoded, err := msg.Encode()
-			if err != nil {
-				errs <- err
-				return
-			}
-			err = ws.WriteMessage(websocket.BinaryMessage, encoded)
-			if utils.IsOKNetworkError(err) {
-				errs <- nil
-				return
-			}
-			if err != nil {
-				errs <- err
-				return
+				return err
 			}
+			tdpMessagesToSend <- msg
 		}
-	}()
+	})
 
-	go func() {
+	// run a goroutine to read TDP messages coming from the browser
+	// and pass them on to the Windows agent
+	errs.Go(func() error {
 		defer closeOnce.Do(close)
 
 		var buf bytes.Buffer
 		for {
 			_, reader, err := ws.NextReader()
 			switch {
 			case utils.IsOKNetworkError(err):
-				errs <- nil
-				return
+				return err
 			case err != nil:
-				errs <- err
-				return
+				return err
 			}
 			buf.Reset()
 			if _, err := io.Copy(&buf, reader); err != nil {
-				errs <- err
-				return
+				return err
 			}
 
 			if _, err := wds.Write(buf.Bytes()); err != nil {
-				errs <- trace.Wrap(err, "sending TDP message to desktop agent")
-				return
+				return trace.Wrap(err, "sending TDP message to desktop agent")
 			}
 		}
-	}()
+	})
 
-	var retErrs []error
-	for i := 0; i < 2; i++ {
-		retErrs = append(retErrs, <-errs)
-	}
-	return trace.NewAggregate(retErrs...)
+	return trace.Wrap(errs.Wait())
 }
 
 // handleProxyWebsocketConnErr handles the error returned by proxyWebsocketConn by