Skip to content

Commit 5b10e59

Browse files
committed
Consolidate decoder validation to DataDecoder
Move size validation logic to DataDecoder to eliminate duplication and create single source of truth for data validation.
1 parent 698918f commit 5b10e59

File tree

4 files changed

+132
-130
lines changed

4 files changed

+132
-130
lines changed

internal/decoder/data_decoder.go

Lines changed: 82 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,12 @@ func (d *DataDecoder) DecodeBytes(size, offset uint) ([]byte, uint, error) {
166166

167167
// DecodeFloat64 decodes a 64-bit float from the given offset.
168168
func (d *DataDecoder) DecodeFloat64(size, offset uint) (float64, uint, error) {
169+
if size != 8 {
170+
return 0, 0, mmdberrors.NewInvalidDatabaseError(
171+
"the MaxMind DB file's data section contains bad data (float 64 size of %v)",
172+
size,
173+
)
174+
}
169175
if offset+size > uint(len(d.buffer)) {
170176
return 0, 0, mmdberrors.NewOffsetError()
171177
}
@@ -177,6 +183,12 @@ func (d *DataDecoder) DecodeFloat64(size, offset uint) (float64, uint, error) {
177183

178184
// DecodeFloat32 decodes a 32-bit float from the given offset.
179185
func (d *DataDecoder) DecodeFloat32(size, offset uint) (float32, uint, error) {
186+
if size != 4 {
187+
return 0, 0, mmdberrors.NewInvalidDatabaseError(
188+
"the MaxMind DB file's data section contains bad data (float32 size of %v)",
189+
size,
190+
)
191+
}
180192
if offset+size > uint(len(d.buffer)) {
181193
return 0, 0, mmdberrors.NewOffsetError()
182194
}
@@ -188,6 +200,12 @@ func (d *DataDecoder) DecodeFloat32(size, offset uint) (float32, uint, error) {
188200

189201
// DecodeInt32 decodes a 32-bit signed integer from the given offset.
190202
func (d *DataDecoder) DecodeInt32(size, offset uint) (int32, uint, error) {
203+
if size > 4 {
204+
return 0, 0, mmdberrors.NewInvalidDatabaseError(
205+
"the MaxMind DB file's data section contains bad data (int32 size of %v)",
206+
size,
207+
)
208+
}
191209
if offset+size > uint(len(d.buffer)) {
192210
return 0, 0, mmdberrors.NewOffsetError()
193211
}
@@ -236,6 +254,18 @@ func (d *DataDecoder) DecodePointer(
236254
return pointer, newOffset, nil
237255
}
238256

257+
// DecodeBool decodes a boolean from the given offset.
258+
func (*DataDecoder) DecodeBool(size, offset uint) (bool, uint, error) {
259+
if size > 1 {
260+
return false, 0, mmdberrors.NewInvalidDatabaseError(
261+
"the MaxMind DB file's data section contains bad data (bool size of %v)",
262+
size,
263+
)
264+
}
265+
value, newOffset := decodeBool(size, offset)
266+
return value, newOffset, nil
267+
}
268+
239269
// DecodeString decodes a string from the given offset.
240270
func (d *DataDecoder) DecodeString(size, offset uint) (string, uint, error) {
241271
if offset+size > uint(len(d.buffer)) {
@@ -249,6 +279,12 @@ func (d *DataDecoder) DecodeString(size, offset uint) (string, uint, error) {
249279

250280
// DecodeUint16 decodes a 16-bit unsigned integer from the given offset.
251281
func (d *DataDecoder) DecodeUint16(size, offset uint) (uint16, uint, error) {
282+
if size > 2 {
283+
return 0, 0, mmdberrors.NewInvalidDatabaseError(
284+
"the MaxMind DB file's data section contains bad data (uint16 size of %v)",
285+
size,
286+
)
287+
}
252288
if offset+size > uint(len(d.buffer)) {
253289
return 0, 0, mmdberrors.NewOffsetError()
254290
}
@@ -265,6 +301,12 @@ func (d *DataDecoder) DecodeUint16(size, offset uint) (uint16, uint, error) {
265301

266302
// DecodeUint32 decodes a 32-bit unsigned integer from the given offset.
267303
func (d *DataDecoder) DecodeUint32(size, offset uint) (uint32, uint, error) {
304+
if size > 4 {
305+
return 0, 0, mmdberrors.NewInvalidDatabaseError(
306+
"the MaxMind DB file's data section contains bad data (uint32 size of %v)",
307+
size,
308+
)
309+
}
268310
if offset+size > uint(len(d.buffer)) {
269311
return 0, 0, mmdberrors.NewOffsetError()
270312
}
@@ -281,6 +323,12 @@ func (d *DataDecoder) DecodeUint32(size, offset uint) (uint32, uint, error) {
281323

282324
// DecodeUint64 decodes a 64-bit unsigned integer from the given offset.
283325
func (d *DataDecoder) DecodeUint64(size, offset uint) (uint64, uint, error) {
326+
if size > 8 {
327+
return 0, 0, mmdberrors.NewInvalidDatabaseError(
328+
"the MaxMind DB file's data section contains bad data (uint64 size of %v)",
329+
size,
330+
)
331+
}
284332
if offset+size > uint(len(d.buffer)) {
285333
return 0, 0, mmdberrors.NewOffsetError()
286334
}
@@ -296,16 +344,32 @@ func (d *DataDecoder) DecodeUint64(size, offset uint) (uint64, uint, error) {
296344
}
297345

298346
// DecodeUint128 decodes a 128-bit unsigned integer from the given offset.
299-
func (d *DataDecoder) DecodeUint128(size, offset uint) (*big.Int, uint, error) {
347+
// Returns the value as high and low 64-bit unsigned integers.
348+
func (d *DataDecoder) DecodeUint128(size, offset uint) (hi, lo uint64, newOffset uint, err error) {
349+
if size > 16 {
350+
return 0, 0, 0, mmdberrors.NewInvalidDatabaseError(
351+
"the MaxMind DB file's data section contains bad data (uint128 size of %v)",
352+
size,
353+
)
354+
}
300355
if offset+size > uint(len(d.buffer)) {
301-
return nil, 0, mmdberrors.NewOffsetError()
356+
return 0, 0, 0, mmdberrors.NewOffsetError()
302357
}
303358

304-
newOffset := offset + size
305-
val := new(big.Int)
306-
val.SetBytes(d.buffer[offset:newOffset])
359+
newOffset = offset + size
307360

308-
return val, newOffset, nil
361+
// Process bytes from most significant to least significant
362+
for _, b := range d.buffer[offset:newOffset] {
363+
var carry byte
364+
lo, carry = append64(lo, b)
365+
hi, _ = append64(hi, carry)
366+
}
367+
368+
return hi, lo, newOffset, nil
369+
}
370+
371+
func append64(val uint64, b byte) (uint64, byte) {
372+
return (val << 8) | uint64(b), byte(val >> 56)
309373
}
310374

311375
// DecodeKey decodes a map key into []byte slice. We use a []byte so that we
@@ -509,12 +573,22 @@ func (d *DataDecoder) decodeFromTypeToDeserializer(
509573

510574
return offset, dser.Uint64(v)
511575
case KindUint128:
512-
v, offset, err := d.DecodeUint128(size, offset)
576+
hi, lo, offset, err := d.DecodeUint128(size, offset)
513577
if err != nil {
514578
return 0, err
515579
}
516580

517-
return offset, dser.Uint128(v)
581+
// Convert hi/lo representation to big.Int for deserializer
582+
value := new(big.Int)
583+
if hi == 0 {
584+
value.SetUint64(lo)
585+
} else {
586+
value.SetUint64(hi)
587+
value.Lsh(value, 64) // Shift high part left by 64 bits
588+
value.Or(value, new(big.Int).SetUint64(lo)) // OR with low part
589+
}
590+
591+
return offset, dser.Uint128(value)
518592
default:
519593
return 0, mmdberrors.NewInvalidDatabaseError("unknown type: %d", dtype)
520594
}

internal/decoder/decoder.go

Lines changed: 16 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -52,28 +52,29 @@ func (d *Decoder) ReadBool() (bool, error) {
5252
return false, d.wrapError(err)
5353
}
5454

55-
if size > 1 {
56-
return false, d.wrapError(mmdberrors.NewInvalidDatabaseError(
57-
"the MaxMind DB file's data section contains bad data (bool size of %v)",
58-
size,
59-
))
55+
value, newOffset, err := d.d.DecodeBool(size, offset)
56+
if err != nil {
57+
return false, d.wrapError(err)
6058
}
61-
62-
var value bool
63-
value, _ = decodeBool(size, offset)
64-
d.setNextOffset(offset)
59+
d.setNextOffset(newOffset)
6560
return value, nil
6661
}
6762

6863
// ReadString reads the value pointed by the decoder as a string.
6964
//
7065
// Returns an error if the database is malformed or if the pointed value is not a string.
7166
func (d *Decoder) ReadString() (string, error) {
72-
val, err := d.readBytes(KindString)
67+
size, offset, err := d.decodeCtrlDataAndFollow(KindString)
7368
if err != nil {
7469
return "", d.wrapError(err)
7570
}
76-
return string(val), nil
71+
72+
value, newOffset, err := d.d.DecodeString(size, offset)
73+
if err != nil {
74+
return "", d.wrapError(err)
75+
}
76+
d.setNextOffset(newOffset)
77+
return value, nil
7778
}
7879

7980
// ReadBytes reads the value pointed by the decoder as bytes.
@@ -96,13 +97,6 @@ func (d *Decoder) ReadFloat32() (float32, error) {
9697
return 0, d.wrapError(err)
9798
}
9899

99-
if size != 4 {
100-
return 0, d.wrapError(mmdberrors.NewInvalidDatabaseError(
101-
"the MaxMind DB file's data section contains bad data (float32 size of %v)",
102-
size,
103-
))
104-
}
105-
106100
value, nextOffset, err := d.d.DecodeFloat32(size, offset)
107101
if err != nil {
108102
return 0, d.wrapError(err)
@@ -121,13 +115,6 @@ func (d *Decoder) ReadFloat64() (float64, error) {
121115
return 0, d.wrapError(err)
122116
}
123117

124-
if size != 8 {
125-
return 0, d.wrapError(mmdberrors.NewInvalidDatabaseError(
126-
"the MaxMind DB file's data section contains bad data (float64 size of %v)",
127-
size,
128-
))
129-
}
130-
131118
value, nextOffset, err := d.d.DecodeFloat64(size, offset)
132119
if err != nil {
133120
return 0, d.wrapError(err)
@@ -146,20 +133,12 @@ func (d *Decoder) ReadInt32() (int32, error) {
146133
return 0, d.wrapError(err)
147134
}
148135

149-
if size > 4 {
150-
return 0, d.wrapError(mmdberrors.NewInvalidDatabaseError(
151-
"the MaxMind DB file's data section contains bad data (int32 size of %v)",
152-
size,
153-
))
154-
}
155-
156136
value, nextOffset, err := d.d.DecodeInt32(size, offset)
157137
if err != nil {
158138
return 0, d.wrapError(err)
159139
}
160140

161141
d.setNextOffset(nextOffset)
162-
163142
return value, nil
164143
}
165144

@@ -172,13 +151,6 @@ func (d *Decoder) ReadUInt16() (uint16, error) {
172151
return 0, d.wrapError(err)
173152
}
174153

175-
if size > 2 {
176-
return 0, d.wrapError(mmdberrors.NewInvalidDatabaseError(
177-
"the MaxMind DB file's data section contains bad data (uint16 size of %v)",
178-
size,
179-
))
180-
}
181-
182154
value, nextOffset, err := d.d.DecodeUint16(size, offset)
183155
if err != nil {
184156
return 0, d.wrapError(err)
@@ -197,13 +169,6 @@ func (d *Decoder) ReadUInt32() (uint32, error) {
197169
return 0, d.wrapError(err)
198170
}
199171

200-
if size > 4 {
201-
return 0, d.wrapError(mmdberrors.NewInvalidDatabaseError(
202-
"the MaxMind DB file's data section contains bad data (uint32 size of %v)",
203-
size,
204-
))
205-
}
206-
207172
value, nextOffset, err := d.d.DecodeUint32(size, offset)
208173
if err != nil {
209174
return 0, d.wrapError(err)
@@ -222,13 +187,6 @@ func (d *Decoder) ReadUInt64() (uint64, error) {
222187
return 0, d.wrapError(err)
223188
}
224189

225-
if size > 8 {
226-
return 0, d.wrapError(mmdberrors.NewInvalidDatabaseError(
227-
"the MaxMind DB file's data section contains bad data (uint64 size of %v)",
228-
size,
229-
))
230-
}
231-
232190
value, nextOffset, err := d.d.DecodeUint64(size, offset)
233191
if err != nil {
234192
return 0, d.wrapError(err)
@@ -247,36 +205,15 @@ func (d *Decoder) ReadUInt128() (hi, lo uint64, err error) {
247205
return 0, 0, d.wrapError(err)
248206
}
249207

250-
if size > 16 {
251-
return 0, 0, d.wrapError(mmdberrors.NewInvalidDatabaseError(
252-
"the MaxMind DB file's data section contains bad data (uint128 size of %v)",
253-
size,
254-
))
255-
}
256-
257-
if offset+size > uint(len(d.d.Buffer())) {
258-
return 0, 0, d.wrapError(mmdberrors.NewInvalidDatabaseError(
259-
"the MaxMind DB file's data section contains bad data (offset+size %d exceeds buffer length %d)",
260-
offset+size,
261-
len(d.d.Buffer()),
262-
))
263-
}
264-
265-
for _, b := range d.d.Buffer()[offset : offset+size] {
266-
var carry byte
267-
lo, carry = append64(lo, b)
268-
hi, _ = append64(hi, carry)
208+
hi, lo, nextOffset, err := d.d.DecodeUint128(size, offset)
209+
if err != nil {
210+
return 0, 0, d.wrapError(err)
269211
}
270212

271-
d.setNextOffset(offset + size)
272-
213+
d.setNextOffset(nextOffset)
273214
return hi, lo, nil
274215
}
275216

276-
func append64(val uint64, b byte) (uint64, byte) {
277-
return (val << 8) | uint64(b), byte(val >> 56)
278-
}
279-
280217
// ReadMap returns an iterator to read the map. The first value from the
281218
// iterator is the key. Please note that this byte slice is only valid during
282219
// the iteration. This is done to avoid an unnecessary allocation. You must

internal/decoder/decoder_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -379,7 +379,7 @@ func TestBoundsChecking(t *testing.T) {
379379
// This should fail gracefully with an error instead of panicking
380380
_, err := decoder.ReadString()
381381
require.Error(t, err)
382-
require.Contains(t, err.Error(), "exceeds buffer length")
382+
require.Contains(t, err.Error(), "unexpected end of database")
383383

384384
// Test DecodeBytes bounds checking with a separate buffer
385385
bytesBuffer := []byte{
@@ -403,7 +403,7 @@ func TestBoundsChecking(t *testing.T) {
403403

404404
_, _, err = decoder2.ReadUInt128()
405405
require.Error(t, err)
406-
require.Contains(t, err.Error(), "exceeds buffer length")
406+
require.Contains(t, err.Error(), "unexpected end of database")
407407
}
408408

409409
func TestPeekKind(t *testing.T) {

0 commit comments

Comments
 (0)