decode.go 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107
  1. package memconservative
  2. import (
  3. "errors"
  4. "fmt"
  5. "io"
  6. "os"
  7. "strings"
  8. "google.golang.org/protobuf/encoding/protowire"
  9. )
  10. var (
  11. errFailedToReadBytes = errors.New("failed to read bytes")
  12. errFailedToReadExpectedLenBytes = errors.New("failed to read expected length of bytes")
  13. errInvalidGeodataFile = errors.New("invalid geodata file")
  14. errInvalidGeodataVarintLength = errors.New("invalid geodata varint length")
  15. errCodeNotFound = errors.New("code not found")
  16. )
  17. func emitBytes(f io.ReadSeeker, code string) ([]byte, error) {
  18. count := 1
  19. isInner := false
  20. tempContainer := make([]byte, 0, 5)
  21. var result []byte
  22. var advancedN uint64 = 1
  23. var geoDataVarintLength, codeVarintLength, varintLenByteLen uint64 = 0, 0, 0
  24. Loop:
  25. for {
  26. container := make([]byte, advancedN)
  27. bytesRead, err := f.Read(container)
  28. if err == io.EOF {
  29. return nil, errCodeNotFound
  30. }
  31. if err != nil {
  32. return nil, errFailedToReadBytes
  33. }
  34. if bytesRead != len(container) {
  35. return nil, errFailedToReadExpectedLenBytes
  36. }
  37. switch count {
  38. case 1, 3: // data type ((field_number << 3) | wire_type)
  39. if container[0] != 10 { // byte `0A` equals to `10` in decimal
  40. return nil, errInvalidGeodataFile
  41. }
  42. advancedN = 1
  43. count++
  44. case 2, 4: // data length
  45. tempContainer = append(tempContainer, container...)
  46. if container[0] > 127 { // max one-byte-length byte `7F`(0FFF FFFF) equals to `127` in decimal
  47. advancedN = 1
  48. goto Loop
  49. }
  50. lenVarint, n := protowire.ConsumeVarint(tempContainer)
  51. if n < 0 {
  52. return nil, errInvalidGeodataVarintLength
  53. }
  54. tempContainer = nil
  55. if !isInner {
  56. isInner = true
  57. geoDataVarintLength = lenVarint
  58. advancedN = 1
  59. } else {
  60. isInner = false
  61. codeVarintLength = lenVarint
  62. varintLenByteLen = uint64(n)
  63. advancedN = codeVarintLength
  64. }
  65. count++
  66. case 5: // data value
  67. if strings.EqualFold(string(container), code) {
  68. count++
  69. offset := -(1 + int64(varintLenByteLen) + int64(codeVarintLength))
  70. _, _ = f.Seek(offset, 1) // back to the start of GeoIP or GeoSite varint
  71. advancedN = geoDataVarintLength // the number of bytes to be read in next round
  72. } else {
  73. count = 1
  74. offset := int64(geoDataVarintLength) - int64(codeVarintLength) - int64(varintLenByteLen) - 1
  75. _, _ = f.Seek(offset, 1) // skip the unmatched GeoIP or GeoSite varint
  76. advancedN = 1 // the next round will be the start of another GeoIPList or GeoSiteList
  77. }
  78. case 6: // matched GeoIP or GeoSite varint
  79. result = container
  80. break Loop
  81. }
  82. }
  83. return result, nil
  84. }
  85. func Decode(filename, code string) ([]byte, error) {
  86. f, err := os.Open(filename)
  87. if err != nil {
  88. return nil, fmt.Errorf("failed to open file: %s, base error: %s", filename, err.Error())
  89. }
  90. defer func(f *os.File) {
  91. _ = f.Close()
  92. }(f)
  93. geoBytes, err := emitBytes(f, code)
  94. if err != nil {
  95. return nil, err
  96. }
  97. return geoBytes, nil
  98. }