WTF is this

common/geosite/geosite_test.go

@@ -0,0 +1,34 @@
+package geosite_test
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/sagernet/sing-box/common/geosite"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestGeosite(t *testing.T) {
+	t.Parallel()
+
+	var buffer bytes.Buffer
+	err := geosite.Write(&buffer, map[string][]geosite.Item{
+		"test": {
+			{
+				Type:  geosite.RuleTypeDomain,
+				Value: "example.org",
+			},
+		},
+	})
+	require.NoError(t, err)
+	reader, codes, err := geosite.NewReader(bytes.NewReader(buffer.Bytes()))
+	require.NoError(t, err)
+	require.Equal(t, []string{"test"}, codes)
+	items, err := reader.Read("test")
+	require.NoError(t, err)
+	require.Equal(t, []geosite.Item{{
+		Type:  geosite.RuleTypeDomain,
+		Value: "example.org",
+	}}, items)
+}

common/geosite/reader.go

@@ -1,17 +1,24 @@
 package geosite
 
 import (
+	"bufio"
+	"encoding/binary"
 	"io"
 	"os"
+	"sync"
+	"sync/atomic"
 
 	E "github.com/sagernet/sing/common/exceptions"
-	"github.com/sagernet/sing/common/rw"
+	"github.com/sagernet/sing/common/varbin"
 )
 
 type Reader struct {
-	reader       io.ReadSeeker
-	domainIndex  map[string]int
-	domainLength map[string]int
+	access         sync.Mutex
+	reader         io.ReadSeeker
+	bufferedReader *bufio.Reader
+	metadataIndex  int64
+	domainIndex    map[string]int
+	domainLength   map[string]int
 }
 
 func Open(path string) (*Reader, []string, error) {
@@ -19,14 +26,22 @@ func Open(path string) (*Reader, []string, error) {
 	if err != nil {
 		return nil, nil, err
 	}
-	reader := &Reader{
-		reader: content,
-	}
-	err = reader.readMetadata()
+	reader, codes, err := NewReader(content)
 	if err != nil {
 		content.Close()
 		return nil, nil, err
 	}
+	return reader, codes, nil
+}
+
+func NewReader(readSeeker io.ReadSeeker) (*Reader, []string, error) {
+	reader := &Reader{
+		reader: readSeeker,
+	}
+	err := reader.readMetadata()
+	if err != nil {
+		return nil, nil, err
+	}
 	codes := make([]string, 0, len(reader.domainIndex))
 	for code := range reader.domainIndex {
 		codes = append(codes, code)
@@ -34,15 +49,23 @@ func Open(path string) (*Reader, []string, error) {
 	return reader, codes, nil
 }
 
+type geositeMetadata struct {
+	Code   string
+	Index  uint64
+	Length uint64
+}
+
 func (r *Reader) readMetadata() error {
-	version, err := rw.ReadByte(r.reader)
+	counter := &readCounter{Reader: r.reader}
+	reader := bufio.NewReader(counter)
+	version, err := reader.ReadByte()
 	if err != nil {
 		return err
 	}
 	if version != 0 {
 		return E.New("unknown version")
 	}
-	entryLength, err := rw.ReadUVariant(r.reader)
+	entryLength, err := binary.ReadUvarint(reader)
 	if err != nil {
 		return err
 	}
@@ -55,16 +78,16 @@ func (r *Reader) readMetadata() error {
 			codeIndex  uint64
 			codeLength uint64
 		)
-		code, err = rw.ReadVString(r.reader)
+		code, err = varbin.ReadValue[string](reader, binary.BigEndian)
 		if err != nil {
 			return err
 		}
 		keys[i] = code
-		codeIndex, err = rw.ReadUVariant(r.reader)
+		codeIndex, err = binary.ReadUvarint(reader)
 		if err != nil {
 			return err
 		}
-		codeLength, err = rw.ReadUVariant(r.reader)
+		codeLength, err = binary.ReadUvarint(reader)
 		if err != nil {
 			return err
 		}
@@ -73,6 +96,8 @@ func (r *Reader) readMetadata() error {
 	}
 	r.domainIndex = domainIndex
 	r.domainLength = domainLength
+	r.metadataIndex = counter.count - int64(reader.Buffered())
+	r.bufferedReader = reader
 	return nil
 }
 
@@ -81,31 +106,32 @@ func (r *Reader) Read(code string) ([]Item, error) {
 	if !exists {
 		return nil, E.New("code ", code, " not exists!")
 	}
-	_, err := r.reader.Seek(int64(index), io.SeekCurrent)
+	_, err := r.reader.Seek(r.metadataIndex+int64(index), io.SeekStart)
 	if err != nil {
 		return nil, err
 	}
-	counter := &rw.ReadCounter{Reader: r.reader}
-	domain := make([]Item, r.domainLength[code])
-	for i := range domain {
-		var (
-			item Item
-			err  error
-		)
-		item.Type, err = rw.ReadByte(counter)
-		if err != nil {
-			return nil, err
-		}
-		item.Value, err = rw.ReadVString(counter)
-		if err != nil {
-			return nil, err
-		}
-		domain[i] = item
+	r.bufferedReader.Reset(r.reader)
+	itemList := make([]Item, r.domainLength[code])
+	err = varbin.Read(r.bufferedReader, binary.BigEndian, &itemList)
+	if err != nil {
+		return nil, err
 	}
-	_, err = r.reader.Seek(int64(-index)-counter.Count(), io.SeekCurrent)
-	return domain, err
+	return itemList, nil
 }
 
 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
+}

common/geosite/writer.go

@@ -2,13 +2,13 @@ package geosite
 
 import (
 	"bytes"
-	"io"
+	"encoding/binary"
 	"sort"
 
-	"github.com/sagernet/sing/common/rw"
+	"github.com/sagernet/sing/common/varbin"
 )
 
-func Write(writer io.Writer, domains map[string][]Item) error {
+func Write(writer varbin.Writer, domains map[string][]Item) error {
 	keys := make([]string, 0, len(domains))
 	for code := range domains {
 		keys = append(keys, code)
@@ -19,35 +19,34 @@ func Write(writer io.Writer, domains map[string][]Item) error {
 	index := make(map[string]int)
 	for _, code := range keys {
 		index[code] = content.Len()
-		for _, domain := range domains[code] {
-			content.WriteByte(domain.Type)
-			err := rw.WriteVString(content, domain.Value)
+		for _, item := range domains[code] {
+			err := varbin.Write(content, binary.BigEndian, item)
 			if err != nil {
 				return err
 			}
 		}
 	}
 
-	err := rw.WriteByte(writer, 0)
+	err := writer.WriteByte(0)
 	if err != nil {
 		return err
 	}
 
-	err = rw.WriteUVariant(writer, uint64(len(keys)))
+	_, err = varbin.WriteUvarint(writer, uint64(len(keys)))
 	if err != nil {
 		return err
 	}
 
 	for _, code := range keys {
-		err = rw.WriteVString(writer, code)
+		err = varbin.Write(writer, binary.BigEndian, code)
 		if err != nil {
 			return err
 		}
-		err = rw.WriteUVariant(writer, uint64(index[code]))
+		_, err = varbin.WriteUvarint(writer, uint64(index[code]))
 		if err != nil {
 			return err
 		}
-		err = rw.WriteUVariant(writer, uint64(len(domains[code])))
+		_, err = varbin.WriteUvarint(writer, uint64(len(domains[code])))
 		if err != nil {
 			return err
 		}