108 lines
2.9 KiB
Go
108 lines
2.9 KiB
Go
|
package memconservative
|
||
|
|
|
||
|
|
import (
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"io"
|
||
|
|
"os"
|
||
|
|
"strings"
|
||
|
|
|
||
|
|
"google.golang.org/protobuf/encoding/protowire"
|
||
|
|
)
|
||
|
|
|
||
|
|
var (
|
||
|
|
errFailedToReadBytes = errors.New("failed to read bytes")
|
||
|
|
errFailedToReadExpectedLenBytes = errors.New("failed to read expected length of bytes")
|
||
|
|
errInvalidGeodataFile = errors.New("invalid geodata file")
|
||
|
|
errInvalidGeodataVarintLength = errors.New("invalid geodata varint length")
|
||
|
|
errCodeNotFound = errors.New("code not found")
|
||
|
|
)
|
||
|
|
|
||
|
|
func emitBytes(f io.ReadSeeker, code string) ([]byte, error) {
|
||
|
|
count := 1
|
||
|
|
isInner := false
|
||
|
|
tempContainer := make([]byte, 0, 5)
|
||
|
|
|
||
|
|
var result []byte
|
||
|
|
var advancedN uint64 = 1
|
||
|
|
var geoDataVarintLength, codeVarintLength, varintLenByteLen uint64 = 0, 0, 0
|
||
|
|
|
||
|
|
Loop:
|
||
|
|
for {
|
||
|
|
container := make([]byte, advancedN)
|
||
|
|
bytesRead, err := f.Read(container)
|
||
|
|
if err == io.EOF {
|
||
|
|
return nil, errCodeNotFound
|
||
|
|
}
|
||
|
|
if err != nil {
|
||
|
|
return nil, errFailedToReadBytes
|
||
|
|
}
|
||
|
|
if bytesRead != len(container) {
|
||
|
|
return nil, errFailedToReadExpectedLenBytes
|
||
|
|
}
|
||
|
|
|
||
|
|
switch count {
|
||
|
|
case 1, 3: // data type ((field_number << 3) | wire_type)
|
||
|
|
if container[0] != 10 { // byte `0A` equals to `10` in decimal
|
||
|
|
return nil, errInvalidGeodataFile
|
||
|
|
}
|
||
|
|
advancedN = 1
|
||
|
|
count++
|
||
|
|
case 2, 4: // data length
|
||
|
|
tempContainer = append(tempContainer, container...)
|
||
|
|
if container[0] > 127 { // max one-byte-length byte `7F`(0FFF FFFF) equals to `127` in decimal
|
||
|
|
advancedN = 1
|
||
|
|
goto Loop
|
||
|
|
}
|
||
|
|
lenVarint, n := protowire.ConsumeVarint(tempContainer)
|
||
|
|
if n < 0 {
|
||
|
|
return nil, errInvalidGeodataVarintLength
|
||
|
|
}
|
||
|
|
tempContainer = nil
|
||
|
|
if !isInner {
|
||
|
|
isInner = true
|
||
|
|
geoDataVarintLength = lenVarint
|
||
|
|
advancedN = 1
|
||
|
|
} else {
|
||
|
|
isInner = false
|
||
|
|
codeVarintLength = lenVarint
|
||
|
|
varintLenByteLen = uint64(n)
|
||
|
|
advancedN = codeVarintLength
|
||
|
|
}
|
||
|
|
count++
|
||
|
|
case 5: // data value
|
||
|
|
if strings.EqualFold(string(container), code) {
|
||
|
|
count++
|
||
|
|
offset := -(1 + int64(varintLenByteLen) + int64(codeVarintLength))
|
||
|
|
_, _ = f.Seek(offset, 1) // back to the start of GeoIP or GeoSite varint
|
||
|
|
advancedN = geoDataVarintLength // the number of bytes to be read in next round
|
||
|
|
} else {
|
||
|
|
count = 1
|
||
|
|
offset := int64(geoDataVarintLength) - int64(codeVarintLength) - int64(varintLenByteLen) - 1
|
||
|
|
_, _ = f.Seek(offset, 1) // skip the unmatched GeoIP or GeoSite varint
|
||
|
|
advancedN = 1 // the next round will be the start of another GeoIPList or GeoSiteList
|
||
|
|
}
|
||
|
|
case 6: // matched GeoIP or GeoSite varint
|
||
|
|
result = container
|
||
|
|
break Loop
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return result, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func Decode(filename, code string) ([]byte, error) {
|
||
|
|
f, err := os.Open(filename)
|
||
|
|
if err != nil {
|
||
|
|
return nil, fmt.Errorf("failed to open file: %s, base error: %s", filename, err.Error())
|
||
|
|
}
|
||
|
|
defer func(f *os.File) {
|
||
|
|
_ = f.Close()
|
||
|
|
}(f)
|
||
|
|
|
||
|
|
geoBytes, err := emitBytes(f, code)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
return geoBytes, nil
|
||
|
|
}
|