sync: accept TDLib zlib gzip_packed payloads

internal/mtprotoedge/encrypted.go

@@ -3,6 +3,7 @@ package mtprotoedge
 import (
 	"bytes"
 	"compress/gzip"
+	"compress/zlib"
 	"context"
 	"crypto/sha256"
 	"encoding/binary"
@@ -413,7 +414,7 @@ func (s *Server) decodeGZIPWithGlobalBudgetLimit(b *bin.Buffer, limit int) ([]by
 		}
 	}
 
-	r, err := gzip.NewReader(bytes.NewReader(compressed))
+	r, err := newGZIPPackedReader(compressed)
 	if err != nil {
 		release()
 		return nil, func() {}, err
@@ -442,6 +443,19 @@ func (s *Server) decodeGZIPWithGlobalBudgetLimit(b *bin.Buffer, limit int) ([]by
 	return data, release, nil
 }
 
+// newGZIPPackedReader accepts the two wrapped DEFLATE formats emitted by
+// official Telegram clients. TDLib uses a zlib wrapper while DrKLO/gotd use a
+// gzip wrapper; raw DEFLATE is deliberately unsupported. Selecting by the gzip
+// magic keeps malformed gzip input on the gzip validator instead of silently
+// retrying it as another format.
+func newGZIPPackedReader(compressed []byte) (io.ReadCloser, error) {
+	source := bytes.NewReader(compressed)
+	if len(compressed) >= 2 && compressed[0] == 0x1f && compressed[1] == 0x8b {
+		return gzip.NewReader(source)
+	}
+	return zlib.NewReader(source)
+}
+
 // gzipPackedBytesView parses the TL bytes envelope without copying the compressed
 // payload. proto.GZIP.Decode calls bin.Buffer.Bytes, which duplicates the compressed
 // frame before allocating the decompressed result.

internal/mtprotoedge/layer_admission_budget_test.go

@@ -776,7 +776,7 @@ func TestLayerRPCAdmissionTransfersOriginalReservationToFreshOwner(t *testing.T)
 func tdlibNestedGZIPBody(t *testing.T, profile tlprofile.Profile, terminal bin.Object) ([]byte, int) {
 	t.Helper()
 	terminalWire := exactOutboundLayerRPCBody(t, profile, terminal)
-	return tdlibWrappedBody(t, profile, &proto.GZIP{Data: terminalWire}), len(terminalWire)
+	return tdlibWrappedBody(t, profile, zlibPackedObjectForTest(t, terminalWire)), len(terminalWire)
 }
 
 func tdlibWrappedBody(t *testing.T, profile tlprofile.Profile, terminal bin.Object) []byte {

internal/mtprotoedge/structural_limits_test.go

@@ -1,6 +1,9 @@
 package mtprotoedge
 
 import (
+	"bytes"
+	"compress/flate"
+	"compress/zlib"
 	"context"
 	"encoding/binary"
 	"errors"
@@ -685,3 +688,121 @@ func TestGZIPExpansionUsesProcessBudgetBeforeDecode(t *testing.T) {
 		t.Fatalf("caller-bounded reservation failure leaked %d bytes", got)
 	}
 }
+
+func TestGZIPPackedAcceptsOfficialClientWrappers(t *testing.T) {
+	payload := []byte("official Telegram gzip_packed payload")
+	tests := []struct {
+		name   string
+		packed bin.Encoder
+	}{
+		{name: "tdlib_zlib", packed: zlibPackedObjectForTest(t, payload)},
+		{name: "drklo_gotd_gzip", packed: &proto.GZIP{Data: payload}},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var wrapped bin.Buffer
+			if err := tt.packed.Encode(&wrapped); err != nil {
+				t.Fatalf("encode gzip_packed: %v", err)
+			}
+
+			s := New(Options{Logger: zaptest.NewLogger(t)})
+			s.frameBudget = newInboundFrameBudget(maxSingleGZIPExpandedBytes)
+			decoded, release, err := s.decodeGZIPWithGlobalBudget(&wrapped)
+			if err != nil {
+				t.Fatalf("decode gzip_packed: %v", err)
+			}
+			if !bytes.Equal(decoded, payload) {
+				t.Fatalf("decoded payload = %q, want %q", decoded, payload)
+			}
+			if got := s.frameBudget.usedBytes(); got != int64(len(payload)) {
+				t.Fatalf("held expansion budget = %d, want %d", got, len(payload))
+			}
+			release()
+			if got := s.frameBudget.usedBytes(); got != 0 {
+				t.Fatalf("released expansion budget = %d, want zero", got)
+			}
+		})
+	}
+}
+
+func TestGZIPPackedRejectsCorruptZLIBChecksumWithoutBudgetLeak(t *testing.T) {
+	packed := zlibPackedObjectForTest(t, []byte("checksum-protected payload"))
+	packed.PackedData[len(packed.PackedData)-1] ^= 0xff
+	var wrapped bin.Buffer
+	if err := packed.Encode(&wrapped); err != nil {
+		t.Fatalf("encode gzip_packed: %v", err)
+	}
+
+	s := New(Options{Logger: zaptest.NewLogger(t)})
+	s.frameBudget = newInboundFrameBudget(maxSingleGZIPExpandedBytes)
+	if _, release, err := s.decodeGZIPWithGlobalBudget(&wrapped); err == nil {
+		release()
+		t.Fatal("corrupt zlib checksum unexpectedly accepted")
+	}
+	if got := s.frameBudget.usedBytes(); got != 0 {
+		t.Fatalf("corrupt zlib checksum leaked %d budget bytes", got)
+	}
+}
+
+func TestGZIPPackedZLIBHonorsExpansionLimit(t *testing.T) {
+	payload := make([]byte, 1<<20)
+	packed := zlibPackedObjectForTest(t, payload)
+	var wrapped bin.Buffer
+	if err := packed.Encode(&wrapped); err != nil {
+		t.Fatalf("encode gzip_packed: %v", err)
+	}
+
+	s := New(Options{Logger: zaptest.NewLogger(t)})
+	s.frameBudget = newInboundFrameBudget(2 * maxSingleGZIPExpandedBytes)
+	if _, release, err := s.decodeGZIPWithGlobalBudgetLimit(&wrapped, len(payload)-1); err == nil {
+		release()
+		t.Fatal("caller-bounded zlib decode accepted an oversized expansion")
+	} else if got := gzipExpansionWork(err); got != len(payload) {
+		t.Fatalf("caller-bounded zlib expansion work = %d, want %d", got, len(payload))
+	}
+	if got := s.frameBudget.usedBytes(); got != 0 {
+		t.Fatalf("caller-bounded zlib rejection leaked %d bytes", got)
+	}
+}
+
+func TestGZIPPackedRejectsRawDEFLATEWithoutBudgetLeak(t *testing.T) {
+	var compressed bytes.Buffer
+	w, err := flate.NewWriter(&compressed, flate.DefaultCompression)
+	if err != nil {
+		t.Fatalf("create raw deflate writer: %v", err)
+	}
+	if _, err := w.Write([]byte("raw deflate is not a supported Telegram wrapper")); err != nil {
+		t.Fatalf("write raw deflate payload: %v", err)
+	}
+	if err := w.Close(); err != nil {
+		t.Fatalf("close raw deflate payload: %v", err)
+	}
+	packed := &mt.GzipPacked{PackedData: compressed.Bytes()}
+	var wrapped bin.Buffer
+	if err := packed.Encode(&wrapped); err != nil {
+		t.Fatalf("encode gzip_packed: %v", err)
+	}
+
+	s := New(Options{Logger: zaptest.NewLogger(t)})
+	s.frameBudget = newInboundFrameBudget(maxSingleGZIPExpandedBytes)
+	if _, release, err := s.decodeGZIPWithGlobalBudget(&wrapped); err == nil {
+		release()
+		t.Fatal("raw deflate unexpectedly accepted")
+	}
+	if got := s.frameBudget.usedBytes(); got != 0 {
+		t.Fatalf("raw deflate rejection leaked %d budget bytes", got)
+	}
+}
+
+func zlibPackedObjectForTest(t testing.TB, payload []byte) *mt.GzipPacked {
+	t.Helper()
+	var compressed bytes.Buffer
+	w := zlib.NewWriter(&compressed)
+	if _, err := w.Write(payload); err != nil {
+		t.Fatalf("write zlib payload: %v", err)
+	}
+	if err := w.Close(); err != nil {
+		t.Fatalf("close zlib payload: %v", err)
+	}
+	return &mt.GzipPacked{PackedData: compressed.Bytes()}
+}