Add LatencyStats decoder (#56483)

lib/srv/desktop/tdp/proto.go

@@ -186,6 +186,8 @@ func decodeMessage(firstByte byte, in byteReader) (Message, error) {
 		return decodeSharedDirectoryTruncateRequest(in)
 	case TypeSharedDirectoryTruncateResponse:
 		return decodeSharedDirectoryTruncateResponse(in)
+	case TypeLatencyStats:
+		return decodeLatencyStats(in)
 	case TypePing:
 		return decodePing(in)
 	case TypeClientKeyboardLayout:
@@ -1653,6 +1655,17 @@ func (l LatencyStats) Encode() ([]byte, error) {
 	return buf.Bytes(), nil
 }
 
+func decodeLatencyStats(in io.Reader) (LatencyStats, error) {
+	var latencyStats LatencyStats
+	if err := binary.Read(in, binary.BigEndian, &latencyStats.ClientLatency); err != nil {
+		return latencyStats, trace.Wrap(err)
+	}
+	if err := binary.Read(in, binary.BigEndian, &latencyStats.ServerLatency); err != nil {
+		return latencyStats, trace.Wrap(err)
+	}
+	return latencyStats, 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 {

lib/web/apiserver_test.go

@@ -2432,6 +2432,12 @@ func TestDesktopAccessMFA(t *testing.T) {
 
 			msg, err := tdpClient.ReadMessage()
 			require.NoError(t, err)
+
+			// sometimes LatencyStats will be sent before we get Alert, in such case just skip it and get next one
+			if _, ok := msg.(tdp.LatencyStats); ok {
+				msg, err = tdpClient.ReadMessage()
+				require.NoError(t, err)
+			}
 			require.IsType(t, tdp.Alert{}, msg)
 		})
 	}