sing-box/common/geosite/reader.go

package geosite

import (
	"bufio"
	"encoding/binary"
	"io"
	"os"
	"sync/atomic"

	"github.com/sagernet/sing/common"
	E "github.com/sagernet/sing/common/exceptions"
	"github.com/sagernet/sing/common/varbin"
)

type Reader struct {
	reader       io.ReadSeeker
	domainIndex  map[string]int
	domainLength map[string]int
}

func Open(path string) (*Reader, []string, error) {
	content, err := os.Open(path)
	if err != nil {
		return nil, nil, err
	}
	reader := &Reader{
		reader: content,
	}
	err = reader.readMetadata()
	if err != nil {
		content.Close()
		return nil, nil, err
	}
	codes := make([]string, 0, len(reader.domainIndex))
	for code := range reader.domainIndex {
		codes = append(codes, code)
	}
	return reader, codes, nil
}

type geositeMetadata struct {
	Code   string
	Index  uint64
	Length uint64
}

func (r *Reader) readMetadata() error {
	reader := bufio.NewReader(r.reader)
	version, err := reader.ReadByte()
	if err != nil {
		return err
	}
	if version != 0 {
		return E.New("unknown version")
	}
	metadataEntries, err := varbin.ReadValue[[]geositeMetadata](reader, binary.BigEndian)
	if err != nil {
		return err
	}
	domainIndex := make(map[string]int)
	domainLength := make(map[string]int)
	for _, entry := range metadataEntries {
		domainIndex[entry.Code] = int(entry.Index)
		domainLength[entry.Code] = int(entry.Length)
	}
	r.domainIndex = domainIndex
	r.domainLength = domainLength
	if reader.Buffered() > 0 {
		return common.Error(r.reader.Seek(int64(-reader.Buffered()), io.SeekCurrent))
	}
	return nil
}

func (r *Reader) Read(code string) ([]Item, error) {
	index, exists := r.domainIndex[code]
	if !exists {
		return nil, E.New("code ", code, " not exists!")
	}
	_, err := r.reader.Seek(int64(index), io.SeekCurrent)
	if err != nil {
		return nil, err
	}
	counter := &readCounter{Reader: r.reader}
	domain, err := varbin.ReadValue[[]Item](bufio.NewReader(counter), binary.BigEndian)
	if err != nil {
		return nil, err
	}
	_, err = r.reader.Seek(int64(-index)-counter.count, io.SeekCurrent)
	return domain, err
}

func (r *Reader) Upstream() any {
	return r.reader
}

type readCounter struct {
	io.Reader
	count int64
}

func (r *readCounter) Read(p []byte) (n int, err error) {
	n, err = r.Reader.Read(p)
	if n > 0 {
		atomic.AddInt64(&r.count, int64(n))
	}
	return
}