github.com/ipfans/trojan-go@v0.11.0/common/geodata/decode.go (about)

     1  // Package geodata includes utilities to decode and parse the geoip & geosite dat files.
     2  //
     3  // It relies on the proto structure of GeoIP, GeoIPList, GeoSite and GeoSiteList in
     4  // github.com/v2fly/v2ray-core/v4/app/router/config.proto to comply with following rules:
     5  //
     6  // 1. GeoIPList and GeoSiteList cannot be changed
     7  // 2. The country_code in GeoIP and GeoSite must be
     8  //    a length-delimited `string`(wired type) and has field_number set to 1
     9  //
    10  package geodata
    11  
    12  import (
    13  	"errors"
    14  	"io"
    15  	"os"
    16  	"strings"
    17  
    18  	"google.golang.org/protobuf/encoding/protowire"
    19  )
    20  
    21  var (
    22  	ErrFailedToReadBytes            = errors.New("failed to read bytes")
    23  	ErrFailedToReadExpectedLenBytes = errors.New("failed to read expected length of bytes")
    24  	ErrInvalidGeodataFile           = errors.New("invalid geodata file")
    25  	ErrInvalidGeodataVarintLength   = errors.New("invalid geodata varint length")
    26  	ErrCodeNotFound                 = errors.New("code not found")
    27  )
    28  
    29  func EmitBytes(f io.ReadSeeker, code string) ([]byte, error) {
    30  	count := 1
    31  	isInner := false
    32  	tempContainer := make([]byte, 0, 5)
    33  
    34  	var result []byte
    35  	var advancedN uint64 = 1
    36  	var geoDataVarintLength, codeVarintLength, varintLenByteLen uint64 = 0, 0, 0
    37  
    38  Loop:
    39  	for {
    40  		container := make([]byte, advancedN)
    41  		bytesRead, err := f.Read(container)
    42  		if err == io.EOF {
    43  			return nil, ErrCodeNotFound
    44  		}
    45  		if err != nil {
    46  			return nil, ErrFailedToReadBytes
    47  		}
    48  		if bytesRead != len(container) {
    49  			return nil, ErrFailedToReadExpectedLenBytes
    50  		}
    51  
    52  		switch count {
    53  		case 1, 3: // data type ((field_number << 3) | wire_type)
    54  			if container[0] != 10 { // byte `0A` equals to `10` in decimal
    55  				return nil, ErrInvalidGeodataFile
    56  			}
    57  			advancedN = 1
    58  			count++
    59  		case 2, 4: // data length
    60  			tempContainer = append(tempContainer, container...)
    61  			if container[0] > 127 { // max one-byte-length byte `7F`(0FFF FFFF) equals to `127` in decimal
    62  				advancedN = 1
    63  				goto Loop
    64  			}
    65  			lenVarint, n := protowire.ConsumeVarint(tempContainer)
    66  			if n < 0 {
    67  				return nil, ErrInvalidGeodataVarintLength
    68  			}
    69  			tempContainer = nil
    70  			if !isInner {
    71  				isInner = true
    72  				geoDataVarintLength = lenVarint
    73  				advancedN = 1
    74  			} else {
    75  				isInner = false
    76  				codeVarintLength = lenVarint
    77  				varintLenByteLen = uint64(n)
    78  				advancedN = codeVarintLength
    79  			}
    80  			count++
    81  		case 5: // data value
    82  			if strings.EqualFold(string(container), code) {
    83  				count++
    84  				offset := -(1 + int64(varintLenByteLen) + int64(codeVarintLength))
    85  				f.Seek(offset, 1)               // back to the start of GeoIP or GeoSite varint
    86  				advancedN = geoDataVarintLength // the number of bytes to be read in next round
    87  			} else {
    88  				count = 1
    89  				offset := int64(geoDataVarintLength) - int64(codeVarintLength) - int64(varintLenByteLen) - 1
    90  				f.Seek(offset, 1) // skip the unmatched GeoIP or GeoSite varint
    91  				advancedN = 1     // the next round will be the start of another GeoIPList or GeoSiteList
    92  			}
    93  		case 6: // matched GeoIP or GeoSite varint
    94  			result = container
    95  			break Loop
    96  		}
    97  	}
    98  
    99  	return result, nil
   100  }
   101  
   102  func Decode(filename, code string) ([]byte, error) {
   103  	f, err := os.Open(filename)
   104  	if err != nil {
   105  		return nil, err
   106  	}
   107  	defer f.Close()
   108  
   109  	geoBytes, err := EmitBytes(f, code)
   110  	if err != nil {
   111  		return nil, err
   112  	}
   113  	return geoBytes, nil
   114  }