decoder.go 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559
  1. package maxminddb
  2. import (
  3. "encoding/binary"
  4. "math"
  5. "math/big"
  6. "reflect"
  7. "sync"
  8. )
  9. type decoder struct {
  10. buffer []byte
  11. }
  12. type dataType int
  13. const (
  14. _Extended dataType = iota
  15. _Pointer
  16. _String
  17. _Float64
  18. _Bytes
  19. _Uint16
  20. _Uint32
  21. _Map
  22. _Int32
  23. _Uint64
  24. _Uint128
  25. _Slice
  26. _Container
  27. _Marker
  28. _Bool
  29. _Float32
  30. )
  31. func (d *decoder) decode(offset uint, result reflect.Value) (uint, error) {
  32. typeNum, size, newOffset := d.decodeCtrlData(offset)
  33. if typeNum != _Pointer && result.Kind() == reflect.Uintptr {
  34. result.Set(reflect.ValueOf(uintptr(offset)))
  35. return d.nextValueOffset(offset, 1), nil
  36. }
  37. return d.decodeFromType(typeNum, size, newOffset, result)
  38. }
  39. func (d *decoder) decodeCtrlData(offset uint) (dataType, uint, uint) {
  40. newOffset := offset + 1
  41. ctrlByte := d.buffer[offset]
  42. typeNum := dataType(ctrlByte >> 5)
  43. if typeNum == _Extended {
  44. typeNum = dataType(d.buffer[newOffset] + 7)
  45. newOffset++
  46. }
  47. var size uint
  48. size, newOffset = d.sizeFromCtrlByte(ctrlByte, newOffset, typeNum)
  49. return typeNum, size, newOffset
  50. }
  51. func (d *decoder) sizeFromCtrlByte(ctrlByte byte, offset uint, typeNum dataType) (uint, uint) {
  52. size := uint(ctrlByte & 0x1f)
  53. if typeNum == _Extended {
  54. return size, offset
  55. }
  56. var bytesToRead uint
  57. if size > 28 {
  58. bytesToRead = size - 28
  59. }
  60. newOffset := offset + bytesToRead
  61. sizeBytes := d.buffer[offset:newOffset]
  62. switch {
  63. case size == 29:
  64. size = 29 + uint(sizeBytes[0])
  65. case size == 30:
  66. size = 285 + uint(uintFromBytes(0, sizeBytes))
  67. case size > 30:
  68. size = uint(uintFromBytes(0, sizeBytes)) + 65821
  69. }
  70. return size, newOffset
  71. }
  72. func (d *decoder) decodeFromType(dtype dataType, size uint, offset uint, result reflect.Value) (uint, error) {
  73. for result.Kind() == reflect.Ptr {
  74. if result.IsNil() {
  75. result.Set(reflect.New(result.Type().Elem()))
  76. }
  77. result = result.Elem()
  78. }
  79. switch dtype {
  80. case _Bool:
  81. return d.unmarshalBool(size, offset, result)
  82. case _Bytes:
  83. return d.unmarshalBytes(size, offset, result)
  84. case _Float32:
  85. return d.unmarshalFloat32(size, offset, result)
  86. case _Float64:
  87. return d.unmarshalFloat64(size, offset, result)
  88. case _Int32:
  89. return d.unmarshalInt32(size, offset, result)
  90. case _Map:
  91. return d.unmarshalMap(size, offset, result)
  92. case _Pointer:
  93. return d.unmarshalPointer(size, offset, result)
  94. case _Slice:
  95. return d.unmarshalSlice(size, offset, result)
  96. case _String:
  97. return d.unmarshalString(size, offset, result)
  98. case _Uint16:
  99. return d.unmarshalUint(size, offset, result, 16)
  100. case _Uint32:
  101. return d.unmarshalUint(size, offset, result, 32)
  102. case _Uint64:
  103. return d.unmarshalUint(size, offset, result, 64)
  104. case _Uint128:
  105. return d.unmarshalUint128(size, offset, result)
  106. default:
  107. return 0, newInvalidDatabaseError("unknown type: %d", dtype)
  108. }
  109. }
  110. func (d *decoder) unmarshalBool(size uint, offset uint, result reflect.Value) (uint, error) {
  111. if size > 1 {
  112. return 0, newInvalidDatabaseError("the MaxMind DB file's data section contains bad data (bool size of %v)", size)
  113. }
  114. value, newOffset, err := d.decodeBool(size, offset)
  115. if err != nil {
  116. return 0, err
  117. }
  118. switch result.Kind() {
  119. default:
  120. return newOffset, newUnmarshalTypeError(value, result.Type())
  121. case reflect.Bool:
  122. result.SetBool(value)
  123. return newOffset, nil
  124. case reflect.Interface:
  125. result.Set(reflect.ValueOf(value))
  126. return newOffset, nil
  127. }
  128. }
  129. func (d *decoder) unmarshalBytes(size uint, offset uint, result reflect.Value) (uint, error) {
  130. value, newOffset, err := d.decodeBytes(size, offset)
  131. if err != nil {
  132. return 0, err
  133. }
  134. switch result.Kind() {
  135. default:
  136. return newOffset, newUnmarshalTypeError(value, result.Type())
  137. case reflect.Slice:
  138. result.SetBytes(value)
  139. return newOffset, nil
  140. case reflect.Interface:
  141. result.Set(reflect.ValueOf(value))
  142. return newOffset, nil
  143. }
  144. }
  145. func (d *decoder) unmarshalFloat32(size uint, offset uint, result reflect.Value) (uint, error) {
  146. if size != 4 {
  147. return 0, newInvalidDatabaseError("the MaxMind DB file's data section contains bad data (float32 size of %v)", size)
  148. }
  149. value, newOffset, err := d.decodeFloat32(size, offset)
  150. if err != nil {
  151. return 0, err
  152. }
  153. switch result.Kind() {
  154. default:
  155. return newOffset, newUnmarshalTypeError(value, result.Type())
  156. case reflect.Float32, reflect.Float64:
  157. result.SetFloat(float64(value))
  158. return newOffset, nil
  159. case reflect.Interface:
  160. result.Set(reflect.ValueOf(value))
  161. return newOffset, nil
  162. }
  163. }
  164. func (d *decoder) unmarshalFloat64(size uint, offset uint, result reflect.Value) (uint, error) {
  165. if size != 8 {
  166. return 0, newInvalidDatabaseError("the MaxMind DB file's data section contains bad data (float 64 size of %v)", size)
  167. }
  168. value, newOffset, err := d.decodeFloat64(size, offset)
  169. if err != nil {
  170. return 0, err
  171. }
  172. switch result.Kind() {
  173. default:
  174. return newOffset, newUnmarshalTypeError(value, result.Type())
  175. case reflect.Float32, reflect.Float64:
  176. if result.OverflowFloat(value) {
  177. return 0, newUnmarshalTypeError(value, result.Type())
  178. }
  179. result.SetFloat(value)
  180. return newOffset, nil
  181. case reflect.Interface:
  182. result.Set(reflect.ValueOf(value))
  183. return newOffset, nil
  184. }
  185. }
  186. func (d *decoder) unmarshalInt32(size uint, offset uint, result reflect.Value) (uint, error) {
  187. if size > 4 {
  188. return 0, newInvalidDatabaseError("the MaxMind DB file's data section contains bad data (int32 size of %v)", size)
  189. }
  190. value, newOffset, err := d.decodeInt(size, offset)
  191. if err != nil {
  192. return 0, err
  193. }
  194. switch result.Kind() {
  195. default:
  196. return newOffset, newUnmarshalTypeError(value, result.Type())
  197. case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
  198. n := int64(value)
  199. if result.OverflowInt(n) {
  200. return 0, newUnmarshalTypeError(value, result.Type())
  201. }
  202. result.SetInt(n)
  203. return newOffset, nil
  204. case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
  205. n := uint64(value)
  206. if result.OverflowUint(n) {
  207. return 0, newUnmarshalTypeError(value, result.Type())
  208. }
  209. result.SetUint(n)
  210. return newOffset, nil
  211. case reflect.Interface:
  212. result.Set(reflect.ValueOf(value))
  213. return newOffset, nil
  214. }
  215. }
  216. func (d *decoder) unmarshalMap(size uint, offset uint, result reflect.Value) (uint, error) {
  217. switch result.Kind() {
  218. default:
  219. return 0, newUnmarshalTypeError("map", result.Type())
  220. case reflect.Struct:
  221. return d.decodeStruct(size, offset, result)
  222. case reflect.Map:
  223. return d.decodeMap(size, offset, result)
  224. case reflect.Interface:
  225. rv := reflect.ValueOf(make(map[string]interface{}, size))
  226. newOffset, err := d.decodeMap(size, offset, rv)
  227. result.Set(rv)
  228. return newOffset, err
  229. }
  230. }
  231. func (d *decoder) unmarshalPointer(size uint, offset uint, result reflect.Value) (uint, error) {
  232. pointer, newOffset := d.decodePointer(size, offset)
  233. _, err := d.decode(pointer, result)
  234. return newOffset, err
  235. }
  236. func (d *decoder) unmarshalSlice(size uint, offset uint, result reflect.Value) (uint, error) {
  237. switch result.Kind() {
  238. default:
  239. return 0, newUnmarshalTypeError("array", result.Type())
  240. case reflect.Slice:
  241. return d.decodeSlice(size, offset, result)
  242. case reflect.Interface:
  243. a := []interface{}{}
  244. rv := reflect.ValueOf(&a).Elem()
  245. newOffset, err := d.decodeSlice(size, offset, rv)
  246. result.Set(rv)
  247. return newOffset, err
  248. }
  249. }
  250. func (d *decoder) unmarshalString(size uint, offset uint, result reflect.Value) (uint, error) {
  251. value, newOffset, err := d.decodeString(size, offset)
  252. if err != nil {
  253. return 0, err
  254. }
  255. switch result.Kind() {
  256. default:
  257. return newOffset, newUnmarshalTypeError(value, result.Type())
  258. case reflect.String:
  259. result.SetString(value)
  260. return newOffset, nil
  261. case reflect.Interface:
  262. result.Set(reflect.ValueOf(value))
  263. return newOffset, nil
  264. }
  265. }
  266. func (d *decoder) unmarshalUint(size uint, offset uint, result reflect.Value, uintType uint) (uint, error) {
  267. if size > uintType/8 {
  268. return 0, newInvalidDatabaseError("the MaxMind DB file's data section contains bad data (uint%v size of %v)", uintType, size)
  269. }
  270. value, newOffset, err := d.decodeUint(size, offset)
  271. if err != nil {
  272. return 0, err
  273. }
  274. switch result.Kind() {
  275. default:
  276. return newOffset, newUnmarshalTypeError(value, result.Type())
  277. case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
  278. n := int64(value)
  279. if result.OverflowInt(n) {
  280. return 0, newUnmarshalTypeError(value, result.Type())
  281. }
  282. result.SetInt(n)
  283. return newOffset, nil
  284. case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
  285. if result.OverflowUint(value) {
  286. return 0, newUnmarshalTypeError(value, result.Type())
  287. }
  288. result.SetUint(value)
  289. return newOffset, nil
  290. case reflect.Interface:
  291. result.Set(reflect.ValueOf(value))
  292. return newOffset, nil
  293. }
  294. }
  295. func (d *decoder) unmarshalUint128(size uint, offset uint, result reflect.Value) (uint, error) {
  296. if size > 16 {
  297. return 0, newInvalidDatabaseError("the MaxMind DB file's data section contains bad data (uint128 size of %v)", size)
  298. }
  299. value, newOffset, err := d.decodeUint128(size, offset)
  300. if err != nil {
  301. return 0, err
  302. }
  303. // XXX - this should allow *big.Int rather than just bigInt
  304. // Currently this is reported as invalid
  305. switch result.Kind() {
  306. default:
  307. return newOffset, newUnmarshalTypeError(value, result.Type())
  308. case reflect.Struct:
  309. result.Set(reflect.ValueOf(*value))
  310. return newOffset, nil
  311. case reflect.Interface, reflect.Ptr:
  312. result.Set(reflect.ValueOf(value))
  313. return newOffset, nil
  314. }
  315. }
  316. func (d *decoder) decodeBool(size uint, offset uint) (bool, uint, error) {
  317. return size != 0, offset, nil
  318. }
  319. func (d *decoder) decodeBytes(size uint, offset uint) ([]byte, uint, error) {
  320. newOffset := offset + size
  321. bytes := make([]byte, size)
  322. copy(bytes, d.buffer[offset:newOffset])
  323. return bytes, newOffset, nil
  324. }
  325. func (d *decoder) decodeFloat64(size uint, offset uint) (float64, uint, error) {
  326. newOffset := offset + size
  327. bits := binary.BigEndian.Uint64(d.buffer[offset:newOffset])
  328. return math.Float64frombits(bits), newOffset, nil
  329. }
  330. func (d *decoder) decodeFloat32(size uint, offset uint) (float32, uint, error) {
  331. newOffset := offset + size
  332. bits := binary.BigEndian.Uint32(d.buffer[offset:newOffset])
  333. return math.Float32frombits(bits), newOffset, nil
  334. }
  335. func (d *decoder) decodeInt(size uint, offset uint) (int, uint, error) {
  336. newOffset := offset + size
  337. var val int32
  338. for _, b := range d.buffer[offset:newOffset] {
  339. val = (val << 8) | int32(b)
  340. }
  341. return int(val), newOffset, nil
  342. }
  343. func (d *decoder) decodeMap(size uint, offset uint, result reflect.Value) (uint, error) {
  344. if result.IsNil() {
  345. result.Set(reflect.MakeMap(result.Type()))
  346. }
  347. for i := uint(0); i < size; i++ {
  348. var key string
  349. var err error
  350. key, offset, err = d.decodeKeyString(offset)
  351. if err != nil {
  352. return 0, err
  353. }
  354. value := reflect.New(result.Type().Elem())
  355. offset, err = d.decode(offset, value)
  356. if err != nil {
  357. return 0, err
  358. }
  359. result.SetMapIndex(reflect.ValueOf(key), value.Elem())
  360. }
  361. return offset, nil
  362. }
  363. func (d *decoder) decodePointer(size uint, offset uint) (uint, uint) {
  364. pointerSize := ((size >> 3) & 0x3) + 1
  365. newOffset := offset + pointerSize
  366. pointerBytes := d.buffer[offset:newOffset]
  367. var prefix uint64
  368. if pointerSize == 4 {
  369. prefix = 0
  370. } else {
  371. prefix = uint64(size & 0x7)
  372. }
  373. unpacked := uint(uintFromBytes(prefix, pointerBytes))
  374. var pointerValueOffset uint
  375. switch pointerSize {
  376. case 1:
  377. pointerValueOffset = 0
  378. case 2:
  379. pointerValueOffset = 2048
  380. case 3:
  381. pointerValueOffset = 526336
  382. case 4:
  383. pointerValueOffset = 0
  384. }
  385. pointer := unpacked + pointerValueOffset
  386. return pointer, newOffset
  387. }
  388. func (d *decoder) decodeSlice(size uint, offset uint, result reflect.Value) (uint, error) {
  389. result.Set(reflect.MakeSlice(result.Type(), int(size), int(size)))
  390. for i := 0; i < int(size); i++ {
  391. var err error
  392. offset, err = d.decode(offset, result.Index(i))
  393. if err != nil {
  394. return 0, err
  395. }
  396. }
  397. return offset, nil
  398. }
  399. func (d *decoder) decodeString(size uint, offset uint) (string, uint, error) {
  400. newOffset := offset + size
  401. return string(d.buffer[offset:newOffset]), newOffset, nil
  402. }
  403. var (
  404. fieldMap = map[reflect.Type]map[string]int{}
  405. fieldMapMu sync.RWMutex
  406. )
  407. func (d *decoder) decodeStruct(size uint, offset uint, result reflect.Value) (uint, error) {
  408. resultType := result.Type()
  409. fieldMapMu.RLock()
  410. fields, ok := fieldMap[resultType]
  411. fieldMapMu.RUnlock()
  412. if !ok {
  413. numFields := resultType.NumField()
  414. fields = make(map[string]int, numFields)
  415. for i := 0; i < numFields; i++ {
  416. fieldType := resultType.Field(i)
  417. fieldName := fieldType.Name
  418. if tag := fieldType.Tag.Get("maxminddb"); tag != "" {
  419. fieldName = tag
  420. }
  421. fields[fieldName] = i
  422. }
  423. fieldMapMu.Lock()
  424. fieldMap[resultType] = fields
  425. fieldMapMu.Unlock()
  426. }
  427. for i := uint(0); i < size; i++ {
  428. var (
  429. err error
  430. key string
  431. )
  432. key, offset, err = d.decodeStructKey(offset)
  433. if err != nil {
  434. return 0, err
  435. }
  436. i, ok := fields[key]
  437. if !ok {
  438. offset = d.nextValueOffset(offset, 1)
  439. continue
  440. }
  441. offset, err = d.decode(offset, result.Field(i))
  442. if err != nil {
  443. return 0, err
  444. }
  445. }
  446. return offset, nil
  447. }
  448. func (d *decoder) decodeUint(size uint, offset uint) (uint64, uint, error) {
  449. newOffset := offset + size
  450. val := uintFromBytes(0, d.buffer[offset:newOffset])
  451. return val, newOffset, nil
  452. }
  453. func (d *decoder) decodeUint128(size uint, offset uint) (*big.Int, uint, error) {
  454. newOffset := offset + size
  455. val := new(big.Int)
  456. val.SetBytes(d.buffer[offset:newOffset])
  457. return val, newOffset, nil
  458. }
  459. func uintFromBytes(prefix uint64, uintBytes []byte) uint64 {
  460. val := prefix
  461. for _, b := range uintBytes {
  462. val = (val << 8) | uint64(b)
  463. }
  464. return val
  465. }
  466. func (d *decoder) decodeKeyString(offset uint) (string, uint, error) {
  467. typeNum, size, newOffset := d.decodeCtrlData(offset)
  468. if typeNum == _Pointer {
  469. pointer, ptrOffset := d.decodePointer(size, newOffset)
  470. key, _, err := d.decodeKeyString(pointer)
  471. return key, ptrOffset, err
  472. }
  473. if typeNum != _String {
  474. return "", 0, newInvalidDatabaseError("unexpected type when decoding string: %v", typeNum)
  475. }
  476. return d.decodeString(size, newOffset)
  477. }
  478. // This function is used to skip ahead to the next value without decoding
  479. // the one at the offset passed in. The size bits have different meanings for
  480. // different data types
  481. func (d *decoder) nextValueOffset(offset uint, numberToSkip uint) uint {
  482. if numberToSkip == 0 {
  483. return offset
  484. }
  485. typeNum, size, offset := d.decodeCtrlData(offset)
  486. switch typeNum {
  487. case _Pointer:
  488. _, offset = d.decodePointer(size, offset)
  489. case _Map:
  490. numberToSkip += 2 * size
  491. case _Slice:
  492. numberToSkip += size
  493. case _Bool:
  494. default:
  495. offset += size
  496. }
  497. return d.nextValueOffset(offset, numberToSkip-1)
  498. }