github.com/sagernet/sing-box@v1.9.0-rc.20/common/geosite/reader.go (about)

     1  package geosite
     2  
     3  import (
     4  	"io"
     5  	"os"
     6  
     7  	E "github.com/sagernet/sing/common/exceptions"
     8  	"github.com/sagernet/sing/common/rw"
     9  )
    10  
    11  type Reader struct {
    12  	reader       io.ReadSeeker
    13  	domainIndex  map[string]int
    14  	domainLength map[string]int
    15  }
    16  
    17  func Open(path string) (*Reader, []string, error) {
    18  	content, err := os.Open(path)
    19  	if err != nil {
    20  		return nil, nil, err
    21  	}
    22  	reader := &Reader{
    23  		reader: content,
    24  	}
    25  	err = reader.readMetadata()
    26  	if err != nil {
    27  		content.Close()
    28  		return nil, nil, err
    29  	}
    30  	codes := make([]string, 0, len(reader.domainIndex))
    31  	for code := range reader.domainIndex {
    32  		codes = append(codes, code)
    33  	}
    34  	return reader, codes, nil
    35  }
    36  
    37  func (r *Reader) readMetadata() error {
    38  	version, err := rw.ReadByte(r.reader)
    39  	if err != nil {
    40  		return err
    41  	}
    42  	if version != 0 {
    43  		return E.New("unknown version")
    44  	}
    45  	entryLength, err := rw.ReadUVariant(r.reader)
    46  	if err != nil {
    47  		return err
    48  	}
    49  	keys := make([]string, entryLength)
    50  	domainIndex := make(map[string]int)
    51  	domainLength := make(map[string]int)
    52  	for i := 0; i < int(entryLength); i++ {
    53  		var (
    54  			code       string
    55  			codeIndex  uint64
    56  			codeLength uint64
    57  		)
    58  		code, err = rw.ReadVString(r.reader)
    59  		if err != nil {
    60  			return err
    61  		}
    62  		keys[i] = code
    63  		codeIndex, err = rw.ReadUVariant(r.reader)
    64  		if err != nil {
    65  			return err
    66  		}
    67  		codeLength, err = rw.ReadUVariant(r.reader)
    68  		if err != nil {
    69  			return err
    70  		}
    71  		domainIndex[code] = int(codeIndex)
    72  		domainLength[code] = int(codeLength)
    73  	}
    74  	r.domainIndex = domainIndex
    75  	r.domainLength = domainLength
    76  	return nil
    77  }
    78  
    79  func (r *Reader) Read(code string) ([]Item, error) {
    80  	index, exists := r.domainIndex[code]
    81  	if !exists {
    82  		return nil, E.New("code ", code, " not exists!")
    83  	}
    84  	_, err := r.reader.Seek(int64(index), io.SeekCurrent)
    85  	if err != nil {
    86  		return nil, err
    87  	}
    88  	counter := &rw.ReadCounter{Reader: r.reader}
    89  	domain := make([]Item, r.domainLength[code])
    90  	for i := range domain {
    91  		var (
    92  			item Item
    93  			err  error
    94  		)
    95  		item.Type, err = rw.ReadByte(counter)
    96  		if err != nil {
    97  			return nil, err
    98  		}
    99  		item.Value, err = rw.ReadVString(counter)
   100  		if err != nil {
   101  			return nil, err
   102  		}
   103  		domain[i] = item
   104  	}
   105  	_, err = r.reader.Seek(int64(-index)-counter.Count(), io.SeekCurrent)
   106  	return domain, err
   107  }
   108  
   109  func (r *Reader) Upstream() any {
   110  	return r.reader
   111  }