From 967e2df8dd16c02a55d610926a3aa747d4eb1212 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Krzysztof=20Skrz=C4=99tnicki?=
 <krzysztof.skrzetnicki@goteleport.com>
Date: Thu, 27 Apr 2023 17:13:21 +0200
Subject: [PATCH] Fuzz TDP protocol, fix two issues. (#25260)

* Add FuzzDecode covering tdp protocol.

* Add FuzzDecode to oss-fuzz

* Limit PNG2Frame size.

* Fix decodeSharedDirectoryAnnounce incorrectly consuming 4 extra bytes.
---
 fuzz/oss-fuzz-build.sh            |  3 +++
 lib/srv/desktop/tdp/proto.go      | 17 +++++++-----
 lib/srv/desktop/tdp/proto_test.go | 44 +++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/fuzz/oss-fuzz-build.sh b/fuzz/oss-fuzz-build.sh
index 355fd81eb25fe..6126ded5ca2a8 100755
--- a/fuzz/oss-fuzz-build.sh
+++ b/fuzz/oss-fuzz-build.sh
@@ -23,6 +23,9 @@ prepare_teleport_api() {
 
 build_teleport_fuzzers() {
 
+  compile_native_go_fuzzer $TELEPORT_PREFIX/lib/srv/desktop/tdp \
+    FuzzDecode fuzz_decode
+
   compile_native_go_fuzzer $TELEPORT_PREFIX/lib/services \
     FuzzParserEvalBoolPredicate fuzz_parser_eval_bool_predicate
 
diff --git a/lib/srv/desktop/tdp/proto.go b/lib/srv/desktop/tdp/proto.go
index 7357f71c90876..fcef7c18772b4 100644
--- a/lib/srv/desktop/tdp/proto.go
+++ b/lib/srv/desktop/tdp/proto.go
@@ -279,6 +279,12 @@ func readRawPNG2Frame(firstByte byte, in byteReader) ([]byte, error) {
 		return nil, trace.Wrap(err)
 	}
 
+	// prevent allocation of giant buffers.
+	// this also avoids panic for due to overflow.
+	if pngLength > maxPNGFrameDataLength {
+		return nil, trace.BadParameter("pngLength too big: %v", pngLength)
+	}
+
 	b := make([]byte, 1+4+pngLength+16)
 	b[0] = firstByte
 
@@ -675,12 +681,8 @@ func (s SharedDirectoryAnnounce) Encode() ([]byte, error) {
 }
 
 func decodeSharedDirectoryAnnounce(in io.Reader) (SharedDirectoryAnnounce, error) {
-	var completionID, directoryID uint32
-	err := binary.Read(in, binary.BigEndian, &completionID)
-	if err != nil {
-		return SharedDirectoryAnnounce{}, trace.Wrap(err)
-	}
-	err = binary.Read(in, binary.BigEndian, &directoryID)
+	var directoryID uint32
+	err := binary.Read(in, binary.BigEndian, &directoryID)
 	if err != nil {
 		return SharedDirectoryAnnounce{}, trace.Wrap(err)
 	}
@@ -1418,6 +1420,9 @@ func writeUint64(b *bytes.Buffer, v uint64) {
 	b.WriteByte(byte(v))
 }
 
+// maxPNGFrameDataLength is maximum data length for PNG2Frame
+const maxPNGFrameDataLength = 10 * 1024 * 1024 // 10MB
+
 // These correspond to TdpErrCode enum in the rust RDP client.
 const (
 	ErrCodeNil           uint32 = 0
diff --git a/lib/srv/desktop/tdp/proto_test.go b/lib/srv/desktop/tdp/proto_test.go
index 07af03d7e2d86..691e9175a599e 100644
--- a/lib/srv/desktop/tdp/proto_test.go
+++ b/lib/srv/desktop/tdp/proto_test.go
@@ -72,6 +72,50 @@ func TestEncodeDecode(t *testing.T) {
 	}
 }
 
+func FuzzDecode(f *testing.F) {
+	var corpus = []string{
+		"0",
+		"\x02",
+		"\x1b\xff\xff\x800",
+		"\x1b\xff\xff\xff\xeb",
+		"\nn\x00\x00\x00\x04  {}",
+		"\v00000000\x00\x00\x00\x00",
+		"\nn\x00\x00\x00\x04 { }000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000",
+	}
+
+	for _, s := range corpus {
+		f.Add([]byte(s))
+	}
+
+	f.Fuzz(func(t *testing.T, buf []byte) {
+		require.NotPanics(t, func() {
+			// decode random buffer
+			msg, err := Decode(buf)
+			if err != nil {
+				return
+			}
+
+			// test that we can encode the message back:
+			buf2, err := msg.Encode()
+			require.NoError(t, err)
+			require.NotNil(t, buf2)
+
+			// decode the new buffer. it must be equal to the original msg.
+			msg2, err := Decode(buf2)
+			require.NoError(t, err)
+			require.Equalf(t, msg, msg2, "mismatch for message %v", buf)
+
+			// encode another time.
+			// after encoding, it must match the second buffer identically.
+			// this isn't the case for the first buffer, as there can be trailing bytes after the message.
+			buf3, err := msg2.Encode()
+			require.NoError(t, err)
+			require.NotNil(t, buf3)
+			require.Equal(t, buf2, buf3)
+		})
+	})
+}
+
 func TestBadDecode(t *testing.T) {
 	// 254 is an unknown message type.
 	_, err := Decode([]byte{254})