WTF is this

2026-04-14 04:38:28 +10:00 · 2024-06-24 09:49:15 +08:00
parent eaa8e2dc70
commit 5efec936da
40 changed files with 376 additions and 506 deletions
--- a/common/geosite/reader.go
+++ b/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) {
@@ -34,15 +41,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 +70,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 +88,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 +98,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
+}
--- a/common/geosite/writer.go
+++ b/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,32 @@ 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)
-			if err != nil {
-				return err
-			}
+		err := varbin.Write(content, binary.BigEndian, domains[code])
+		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
 		}
--- a/common/srs/binary.go
+++ b/common/srs/binary.go
@@ -1,6 +1,7 @@
 package srs

 import (
+	"bufio"
 	"compress/zlib"
 	"encoding/binary"
 	"io"
@@ -11,7 +12,7 @@ import (
 	"github.com/sagernet/sing/common"
 	"github.com/sagernet/sing/common/domain"
 	E "github.com/sagernet/sing/common/exceptions"
-	"github.com/sagernet/sing/common/rw"
+	"github.com/sagernet/sing/common/varbin"

 	"go4.org/netipx"
 )
@@ -38,7 +39,7 @@ const (
 	ruleItemFinal uint8 = 0xFF
 )

-func Read(reader io.Reader, recovery bool) (ruleSet option.PlainRuleSet, err error) {
+func Read(reader io.Reader, recover bool) (ruleSet option.PlainRuleSet, err error) {
 	var magicBytes [3]byte
 	_, err = io.ReadFull(reader, magicBytes[:])
 	if err != nil {
@@ -60,13 +61,14 @@ func Read(reader io.Reader, recovery bool) (ruleSet option.PlainRuleSet, err err
 	if err != nil {
 		return
 	}
-	length, err := rw.ReadUVariant(zReader)
+	bReader := bufio.NewReader(zReader)
+	length, err := binary.ReadUvarint(bReader)
 	if err != nil {
 		return
 	}
 	ruleSet.Rules = make([]option.HeadlessRule, length)
 	for i := uint64(0); i < length; i++ {
-		ruleSet.Rules[i], err = readRule(zReader, recovery)
+		ruleSet.Rules[i], err = readRule(bReader, recover)
 		if err != nil {
 			err = E.Cause(err, "read rule[", i, "]")
 			return
@@ -88,20 +90,25 @@ func Write(writer io.Writer, ruleSet option.PlainRuleSet) error {
 	if err != nil {
 		return err
 	}
-	err = rw.WriteUVariant(zWriter, uint64(len(ruleSet.Rules)))
+	bWriter := bufio.NewWriter(zWriter)
+	_, err = varbin.WriteUvarint(bWriter, uint64(len(ruleSet.Rules)))
 	if err != nil {
 		return err
 	}
 	for _, rule := range ruleSet.Rules {
-		err = writeRule(zWriter, rule)
+		err = writeRule(bWriter, rule)
 		if err != nil {
 			return err
 		}
 	}
+	err = bWriter.Flush()
+	if err != nil {
+		return err
+	}
 	return zWriter.Close()
 }

-func readRule(reader io.Reader, recovery bool) (rule option.HeadlessRule, err error) {
+func readRule(reader varbin.Reader, recover bool) (rule option.HeadlessRule, err error) {
 	var ruleType uint8
 	err = binary.Read(reader, binary.BigEndian, &ruleType)
 	if err != nil {
@@ -110,17 +117,17 @@ func readRule(reader io.Reader, recovery bool) (rule option.HeadlessRule, err er
 	switch ruleType {
 	case 0:
 		rule.Type = C.RuleTypeDefault
-		rule.DefaultOptions, err = readDefaultRule(reader, recovery)
+		rule.DefaultOptions, err = readDefaultRule(reader, recover)
 	case 1:
 		rule.Type = C.RuleTypeLogical
-		rule.LogicalOptions, err = readLogicalRule(reader, recovery)
+		rule.LogicalOptions, err = readLogicalRule(reader, recover)
 	default:
 		err = E.New("unknown rule type: ", ruleType)
 	}
 	return
 }

-func writeRule(writer io.Writer, rule option.HeadlessRule) error {
+func writeRule(writer varbin.Writer, rule option.HeadlessRule) error {
 	switch rule.Type {
 	case C.RuleTypeDefault:
 		return writeDefaultRule(writer, rule.DefaultOptions)
@@ -131,7 +138,7 @@ func writeRule(writer io.Writer, rule option.HeadlessRule) error {
 	}
 }

-func readDefaultRule(reader io.Reader, recovery bool) (rule option.DefaultHeadlessRule, err error) {
+func readDefaultRule(reader varbin.Reader, recover bool) (rule option.DefaultHeadlessRule, err error) {
 	var lastItemType uint8
 	for {
 		var itemType uint8
@@ -158,6 +165,9 @@ func readDefaultRule(reader io.Reader, recovery bool) (rule option.DefaultHeadle
 				return
 			}
 			rule.DomainMatcher = matcher
+			if recover {
+				rule.Domain, rule.DomainSuffix = matcher.Dump()
+			}
 		case ruleItemDomainKeyword:
 			rule.DomainKeyword, err = readRuleItemString(reader)
 		case ruleItemDomainRegex:
@@ -167,7 +177,7 @@ func readDefaultRule(reader io.Reader, recovery bool) (rule option.DefaultHeadle
 			if err != nil {
 				return
 			}
-			if recovery {
+			if recover {
 				rule.SourceIPCIDR = common.Map(rule.SourceIPSet.Prefixes(), netip.Prefix.String)
 			}
 		case ruleItemIPCIDR:
@@ -175,7 +185,7 @@ func readDefaultRule(reader io.Reader, recovery bool) (rule option.DefaultHeadle
 			if err != nil {
 				return
 			}
-			if recovery {
+			if recover {
 				rule.IPCIDR = common.Map(rule.IPSet.Prefixes(), netip.Prefix.String)
 			}
 		case ruleItemSourcePort:
@@ -209,7 +219,7 @@ func readDefaultRule(reader io.Reader, recovery bool) (rule option.DefaultHeadle
 	}
 }

-func writeDefaultRule(writer io.Writer, rule option.DefaultHeadlessRule) error {
+func writeDefaultRule(writer varbin.Writer, rule option.DefaultHeadlessRule) error {
 	err := binary.Write(writer, binary.BigEndian, uint8(0))
 	if err != nil {
 		return err
@@ -327,73 +337,31 @@ func writeDefaultRule(writer io.Writer, rule option.DefaultHeadlessRule) error {
 	return nil
 }

-func readRuleItemString(reader io.Reader) ([]string, error) {
-	length, err := rw.ReadUVariant(reader)
-	if err != nil {
-		return nil, err
-	}
-	value := make([]string, length)
-	for i := uint64(0); i < length; i++ {
-		value[i], err = rw.ReadVString(reader)
-		if err != nil {
-			return nil, err
-		}
-	}
-	return value, nil
+func readRuleItemString(reader varbin.Reader) ([]string, error) {
+	return varbin.ReadValue[[]string](reader, binary.BigEndian)
 }

-func writeRuleItemString(writer io.Writer, itemType uint8, value []string) error {
-	err := binary.Write(writer, binary.BigEndian, itemType)
+func writeRuleItemString(writer varbin.Writer, itemType uint8, value []string) error {
+	err := writer.WriteByte(itemType)
 	if err != nil {
 		return err
 	}
-	err = rw.WriteUVariant(writer, uint64(len(value)))
+	return varbin.Write(writer, binary.BigEndian, value)
+}
+
+func readRuleItemUint16(reader varbin.Reader) ([]uint16, error) {
+	return varbin.ReadValue[[]uint16](reader, binary.BigEndian)
+}
+
+func writeRuleItemUint16(writer varbin.Writer, itemType uint8, value []uint16) error {
+	err := writer.WriteByte(itemType)
 	if err != nil {
 		return err
 	}
-	for _, item := range value {
-		err = rw.WriteVString(writer, item)
-		if err != nil {
-			return err
-		}
-	}
-	return nil
+	return varbin.Write(writer, binary.BigEndian, value)
 }

-func readRuleItemUint16(reader io.Reader) ([]uint16, error) {
-	length, err := rw.ReadUVariant(reader)
-	if err != nil {
-		return nil, err
-	}
-	value := make([]uint16, length)
-	for i := uint64(0); i < length; i++ {
-		err = binary.Read(reader, binary.BigEndian, &value[i])
-		if err != nil {
-			return nil, err
-		}
-	}
-	return value, nil
-}
-
-func writeRuleItemUint16(writer io.Writer, itemType uint8, value []uint16) error {
-	err := binary.Write(writer, binary.BigEndian, itemType)
-	if err != nil {
-		return err
-	}
-	err = rw.WriteUVariant(writer, uint64(len(value)))
-	if err != nil {
-		return err
-	}
-	for _, item := range value {
-		err = binary.Write(writer, binary.BigEndian, item)
-		if err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
-func writeRuleItemCIDR(writer io.Writer, itemType uint8, value []string) error {
+func writeRuleItemCIDR(writer varbin.Writer, itemType uint8, value []string) error {
 	var builder netipx.IPSetBuilder
 	for i, prefixString := range value {
 		prefix, err := netip.ParsePrefix(prefixString)
@@ -419,9 +387,8 @@ func writeRuleItemCIDR(writer io.Writer, itemType uint8, value []string) error {
 	return writeIPSet(writer, ipSet)
 }

-func readLogicalRule(reader io.Reader, recovery bool) (logicalRule option.LogicalHeadlessRule, err error) {
-	var mode uint8
-	err = binary.Read(reader, binary.BigEndian, &mode)
+func readLogicalRule(reader varbin.Reader, recovery bool) (logicalRule option.LogicalHeadlessRule, err error) {
+	mode, err := reader.ReadByte()
 	if err != nil {
 		return
 	}
@@ -434,7 +401,7 @@ func readLogicalRule(reader io.Reader, recovery bool) (logicalRule option.Logica
 		err = E.New("unknown logical mode: ", mode)
 		return
 	}
-	length, err := rw.ReadUVariant(reader)
+	length, err := binary.ReadUvarint(reader)
 	if err != nil {
 		return
 	}
@@ -453,7 +420,7 @@ func readLogicalRule(reader io.Reader, recovery bool) (logicalRule option.Logica
 	return
 }

-func writeLogicalRule(writer io.Writer, logicalRule option.LogicalHeadlessRule) error {
+func writeLogicalRule(writer varbin.Writer, logicalRule option.LogicalHeadlessRule) error {
 	err := binary.Write(writer, binary.BigEndian, uint8(1))
 	if err != nil {
 		return err
@@ -469,7 +436,7 @@ func writeLogicalRule(writer io.Writer, logicalRule option.LogicalHeadlessRule)
 	if err != nil {
 		return err
 	}
-	err = rw.WriteUVariant(writer, uint64(len(logicalRule.Rules)))
+	_, err = varbin.WriteUvarint(writer, uint64(len(logicalRule.Rules)))
 	if err != nil {
 		return err
 	}
--- a/common/srs/ip_set.go
+++ b/common/srs/ip_set.go
@@ -2,11 +2,13 @@ package srs

 import (
 	"encoding/binary"
-	"io"
 	"net/netip"
+	"os"
 	"unsafe"

-	"github.com/sagernet/sing/common/rw"
+	"github.com/sagernet/sing/common"
+	M "github.com/sagernet/sing/common/metadata"
+	"github.com/sagernet/sing/common/varbin"

 	"go4.org/netipx"
 )
@@ -20,94 +22,57 @@ type myIPRange struct {
 	to   netip.Addr
 }

-func readIPSet(reader io.Reader) (*netipx.IPSet, error) {
-	var version uint8
-	err := binary.Read(reader, binary.BigEndian, &version)
+type myIPRangeData struct {
+	From []byte
+	To   []byte
+}
+
+func readIPSet(reader varbin.Reader) (*netipx.IPSet, error) {
+	version, err := reader.ReadByte()
 	if err != nil {
 		return nil, err
 	}
+	if version != 1 {
+		return nil, os.ErrInvalid
+	}
+	// WTF why using uint64 here
 	var length uint64
 	err = binary.Read(reader, binary.BigEndian, &length)
 	if err != nil {
 		return nil, err
 	}
-	mySet := &myIPSet{
-		rr: make([]myIPRange, length),
+	ranges := make([]myIPRangeData, length)
+	err = varbin.Read(reader, binary.BigEndian, &ranges)
+	if err != nil {
+		return nil, err
 	}
-	for i := uint64(0); i < length; i++ {
-		var (
-			fromLen  uint64
-			toLen    uint64
-			fromAddr netip.Addr
-			toAddr   netip.Addr
-		)
-		fromLen, err = rw.ReadUVariant(reader)
-		if err != nil {
-			return nil, err
-		}
-		fromBytes := make([]byte, fromLen)
-		_, err = io.ReadFull(reader, fromBytes)
-		if err != nil {
-			return nil, err
-		}
-		err = fromAddr.UnmarshalBinary(fromBytes)
-		if err != nil {
-			return nil, err
-		}
-		toLen, err = rw.ReadUVariant(reader)
-		if err != nil {
-			return nil, err
-		}
-		toBytes := make([]byte, toLen)
-		_, err = io.ReadFull(reader, toBytes)
-		if err != nil {
-			return nil, err
-		}
-		err = toAddr.UnmarshalBinary(toBytes)
-		if err != nil {
-			return nil, err
-		}
-		mySet.rr[i] = myIPRange{fromAddr, toAddr}
+	mySet := &myIPSet{
+		rr: make([]myIPRange, len(ranges)),
+	}
+	for i, rangeData := range ranges {
+		mySet.rr[i].from = M.AddrFromIP(rangeData.From)
+		mySet.rr[i].to = M.AddrFromIP(rangeData.To)
 	}
 	return (*netipx.IPSet)(unsafe.Pointer(mySet)), nil
 }

-func writeIPSet(writer io.Writer, set *netipx.IPSet) error {
-	err := binary.Write(writer, binary.BigEndian, uint8(1))
+func writeIPSet(writer varbin.Writer, set *netipx.IPSet) error {
+	err := writer.WriteByte(1)
 	if err != nil {
 		return err
 	}
-	mySet := (*myIPSet)(unsafe.Pointer(set))
-	err = binary.Write(writer, binary.BigEndian, uint64(len(mySet.rr)))
+	dataList := common.Map((*myIPSet)(unsafe.Pointer(set)).rr, func(rr myIPRange) myIPRangeData {
+		return myIPRangeData{
+			From: rr.from.AsSlice(),
+			To:   rr.to.AsSlice(),
+		}
+	})
+	err = binary.Write(writer, binary.BigEndian, uint64(len(dataList)))
 	if err != nil {
 		return err
 	}
-	for _, rr := range mySet.rr {
-		var (
-			fromBinary []byte
-			toBinary   []byte
-		)
-		fromBinary, err = rr.from.MarshalBinary()
-		if err != nil {
-			return err
-		}
-		err = rw.WriteUVariant(writer, uint64(len(fromBinary)))
-		if err != nil {
-			return err
-		}
-		_, err = writer.Write(fromBinary)
-		if err != nil {
-			return err
-		}
-		toBinary, err = rr.to.MarshalBinary()
-		if err != nil {
-			return err
-		}
-		err = rw.WriteUVariant(writer, uint64(len(toBinary)))
-		if err != nil {
-			return err
-		}
-		_, err = writer.Write(toBinary)
+	for _, data := range dataList {
+		err = varbin.Write(writer, binary.BigEndian, data)
 		if err != nil {
 			return err
 		}