diff --git a/decode_map.go b/decode_map.go index 0900c16d..51dd6793 100644 --- a/decode_map.go +++ b/decode_map.go @@ -250,7 +250,12 @@ func (d *Decoder) DecodeUntypedMap() (map[interface{}]interface{}, error) { return nil, nil } - m := make(map[interface{}]interface{}, n) + ln := n + if d.flags&disableAllocLimitFlag == 0 { + ln = min(ln, maxMapSize) + } + + m := make(map[interface{}]interface{}, ln) for i := 0; i < n; i++ { mk, err := d.decodeInterfaceCond() diff --git a/msgpack_test.go b/msgpack_test.go index 296adf8c..d99abb17 100644 --- a/msgpack_test.go +++ b/msgpack_test.go @@ -6,6 +6,7 @@ import ( "fmt" "math" "reflect" + "runtime" "testing" "time" @@ -75,6 +76,45 @@ func (t *MsgpackTest) TestLargeString() { t.Equal(dst, src) } +func (t *MsgpackTest) TestDecodeUntypedMapHugeDeclaredLen() { + // A map32 header declaring ~2G entries with no payload: the map size + // hint must be clamped at maxMapSize before allocation (with the old + // code this allocated a multi-GB map upfront), then fail decoding the + // first key. + // + // 0x7fffffff rather than 0xffffffff: the latter is now rejected by the + // int-overflow check on 32-bit builds before reaching the clamp, so it + // would pass for the wrong reason. This value fits in an int on every + // platform and exercises the clamp itself. + data := []byte{0xdf, 0x7f, 0xff, 0xff, 0xff} + + var before, after runtime.MemStats + runtime.GC() + runtime.ReadMemStats(&before) + + dec := msgpack.NewDecoder(bytes.NewReader(data)) + _, err := dec.DecodeUntypedMap() + t.NotNil(err) + + runtime.ReadMemStats(&after) + // The clamp caps the hint at maxMapSize (1M entries). Without it the + // declared ~2G entries are allocated upfront -- hundreds of MB and tens + // of seconds. Allow generous headroom while still failing loudly if the + // clamp is removed. + if alloc := after.TotalAlloc - before.TotalAlloc; alloc > 256<<20 { + t.Failf("allocation not clamped", "decode allocated %d MiB, want < 256 MiB", alloc>>20) + } +} + +func (t *MsgpackTest) TestDecodeUntypedMap() { + in := map[interface{}]interface{}{int8(1): "one", "two": int8(2)} + t.Nil(t.enc.Encode(in)) + + out, err := t.dec.DecodeUntypedMap() + t.Nil(err) + t.Equal(in, out) +} + func (t *MsgpackTest) TestSliceOfStructs() { in := []*nameStruct{{"hello"}} var out []*nameStruct