diff --git a/lib/go/thrift/simple_json_protocol.go b/lib/go/thrift/simple_json_protocol.go index ec12991a1d..6ae3cba566 100644 --- a/lib/go/thrift/simple_json_protocol.go +++ b/lib/go/thrift/simple_json_protocol.go @@ -432,7 +432,13 @@ func (p *TSimpleJSONProtocol) ReadMapBegin(ctx context.Context) (keyType TType, if err != nil { return keyType, valueType, 0, err } - err = checkSizeForProtocol(int32(size), p.cfg) + if iSize > math.MaxInt32 { + return keyType, valueType, 0, NewTProtocolExceptionWithType( + SIZE_LIMIT, + fmt.Errorf("size exceeded max allowed: %d", iSize), + ) + } + err = checkSizeForProtocol(int32(iSize), p.cfg) if err != nil { return keyType, valueType, 0, err } diff --git a/lib/go/thrift/simple_json_protocol_test.go b/lib/go/thrift/simple_json_protocol_test.go index 18650989a0..e7f2ac95a7 100644 --- a/lib/go/thrift/simple_json_protocol_test.go +++ b/lib/go/thrift/simple_json_protocol_test.go @@ -791,3 +791,44 @@ func TestJSONContextStack(t *testing.T) { func TestTSimpleJSONProtocolUnmatchedBeginEnd(t *testing.T) { UnmatchedBeginEndProtocolTest(t, NewTSimpleJSONProtocolFactory()) } + +func TestReadSimpleJSONProtocolMapBeginSizeLimit(t *testing.T) { + ctx := context.Background() + trans := NewTMemoryBuffer() + wp := NewTSimpleJSONProtocol(trans) + if err := wp.WriteMapBegin(ctx, STRING, STRING, 1<<30); err != nil { + t.Fatal(err) + } + if err := wp.Flush(ctx); err != nil { + t.Fatal(err) + } + + rp := NewTSimpleJSONProtocolConf(trans, &TConfiguration{MaxMessageSize: 1024}) + _, _, size, err := rp.ReadMapBegin(ctx) + if err == nil { + t.Fatalf("expected size-limit error reading oversized map, got nil with size %d", size) + } + if terr, ok := err.(TProtocolException); !ok || terr.TypeId() != SIZE_LIMIT { + t.Errorf("expected SIZE_LIMIT protocol exception, got %v", err) + } +} + +func TestReadSimpleJSONProtocolMapBeginSizeOverflow(t *testing.T) { + // iSize = 1<<32 + 1; int32 narrowing wraps it to 1, which would pass checkSizeForProtocol. + // The map header is written directly as raw JSON to avoid int-width issues on 32-bit platforms. + overflowSize := int64(math.MaxInt32)*2 + 3 // = 4294967297 = 1<<32 + 1 + buf := NewTMemoryBuffer() + buf.WriteString(fmt.Sprintf("[%d,%d,%d]", I32, I32, overflowSize)) + proto := NewTSimpleJSONProtocolConf(buf, &TConfiguration{MaxMessageSize: 1024}) + _, _, _, err := proto.ReadMapBegin(context.Background()) + if err == nil { + t.Fatal("expected error, got nil") + } + tpe, ok := err.(TProtocolException) + if !ok { + t.Fatalf("expected TProtocolException, got %T: %v", err, err) + } + if tpe.TypeId() != SIZE_LIMIT { + t.Errorf("expected SIZE_LIMIT, got %d: %v", tpe.TypeId(), err) + } +}