sing-box/common/srs/binary.go

package srs

import (
	"bufio"
	"compress/zlib"
	"encoding/binary"
	"io"
	"net/netip"
	"slices"
	"unsafe"

	"github.com/sagernet/sing-box/common/ipset"
	C "github.com/sagernet/sing-box/constant"
	"github.com/sagernet/sing-box/option"
	"github.com/sagernet/sing/common"
	"github.com/sagernet/sing/common/domain"
	E "github.com/sagernet/sing/common/exceptions"
	"github.com/sagernet/sing/common/json/badjson"
	"github.com/sagernet/sing/common/json/badoption"
	"github.com/sagernet/sing/common/varbin"

	"go4.org/netipx"
)

var MagicBytes = [3]byte{0x53, 0x52, 0x53} // SRS

const (
	ruleItemQueryType uint8 = iota
	ruleItemNetwork
	ruleItemDomain
	ruleItemDomainKeyword
	ruleItemDomainRegex
	ruleItemSourceIPCIDR
	ruleItemIPCIDR
	ruleItemSourcePort
	ruleItemSourcePortRange
	ruleItemPort
	ruleItemPortRange
	ruleItemProcessName
	ruleItemProcessPath
	ruleItemPackageName
	ruleItemWIFISSID
	ruleItemWIFIBSSID
	ruleItemAdGuardDomain
	ruleItemProcessPathRegex
	ruleItemNetworkType
	ruleItemNetworkIsExpensive
	ruleItemNetworkIsConstrained
	ruleItemNetworkInterfaceAddress
	ruleItemDefaultInterfaceAddress
	ruleItemPackageNameRegex
	ruleItemFinal uint8 = 0xFF
)

func Read(reader io.Reader, recover bool) (ruleSetCompat option.PlainRuleSetCompat, err error) {
	var magicBytes [3]byte
	_, err = io.ReadFull(reader, magicBytes[:])
	if err != nil {
		return
	}
	if magicBytes != MagicBytes {
		err = E.New("invalid sing-box rule-set file")
		return
	}
	var version uint8
	err = binary.Read(reader, binary.BigEndian, &version)
	if err != nil {
		return ruleSetCompat, err
	}
	if version > C.RuleSetVersionCurrent {
		return ruleSetCompat, E.New("unsupported version: ", version)
	}
	compressReader, err := zlib.NewReader(reader)
	if err != nil {
		return
	}
	bReader := bufio.NewReader(compressReader)
	length, err := binary.ReadUvarint(bReader)
	if err != nil {
		return
	}
	ruleSetCompat.Version = version
	for i := range length {
		var rule option.HeadlessRule
		rule, err = readRule(bReader, recover, 0, nil)
		if err != nil {
			err = E.Cause(err, "read rule[", i, "]")
			return
		}
		ruleSetCompat.Options.Rules = append(ruleSetCompat.Options.Rules, rule)
	}
	return
}

func Write(writer io.Writer, ruleSet option.PlainRuleSet, generateVersion uint8) error {
	_, err := writer.Write(MagicBytes[:])
	if err != nil {
		return err
	}
	err = binary.Write(writer, binary.BigEndian, generateVersion)
	if err != nil {
		return err
	}
	compressWriter, err := zlib.NewWriterLevel(writer, zlib.BestCompression)
	if err != nil {
		return err
	}
	bWriter := bufio.NewWriter(compressWriter)
	_, err = varbin.WriteUvarint(bWriter, uint64(len(ruleSet.Rules)))
	if err != nil {
		return err
	}
	for _, rule := range ruleSet.Rules {
		err = writeRule(bWriter, rule, generateVersion, nil)
		if err != nil {
			return err
		}
	}
	err = bWriter.Flush()
	if err != nil {
		return err
	}
	return compressWriter.Close()
}

const maxLogicalRuleDepth = 100

func readRule(reader varbin.Reader, recover bool, depth int, mmap *mmapReader) (rule option.HeadlessRule, err error) {
	if depth > maxLogicalRuleDepth {
		err = E.New("logical rule nested too deep")
		return
	}
	var ruleType uint8
	err = binary.Read(reader, binary.BigEndian, &ruleType)
	if err != nil {
		return
	}
	switch ruleType {
	case 0:
		rule.Type = C.RuleTypeDefault
		rule.DefaultOptions, err = readDefaultRule(reader, recover, mmap)
	case 1:
		rule.Type = C.RuleTypeLogical
		rule.LogicalOptions, err = readLogicalRule(reader, recover, depth, mmap)
	default:
		err = E.New("unknown rule type: ", ruleType)
	}
	return
}

func writeRule(writer varbin.Writer, rule option.HeadlessRule, generateVersion uint8, mmap *mmapWriter) error {
	switch rule.Type {
	case C.RuleTypeDefault:
		return writeDefaultRule(writer, rule.DefaultOptions, generateVersion, mmap)
	case C.RuleTypeLogical:
		return writeLogicalRule(writer, rule.LogicalOptions, generateVersion, mmap)
	default:
		return E.New("unknown rule type: ", rule.Type)
	}
}

func readDefaultRule(reader varbin.Reader, recover bool, mmap *mmapReader) (rule option.DefaultHeadlessRule, err error) {
	var lastItemType uint8
	for {
		var itemType uint8
		err = binary.Read(reader, binary.BigEndian, &itemType)
		if err != nil {
			return
		}
		switch itemType {
		case ruleItemQueryType:
			var rawQueryType []uint16
			rawQueryType, err = varbin.ReadSlice[uint16](reader, binary.BigEndian)
			if err != nil {
				return
			}
			rule.QueryType = common.Map(rawQueryType, func(it uint16) option.DNSQueryType {
				return option.DNSQueryType(it)
			})
		case ruleItemNetwork:
			rule.Network, err = readRuleItemString(reader)
		case ruleItemDomain:
			var matcher *domain.Matcher
			if mmap != nil {
				matcher, err = mmap.readMatcher(reader)
			} else {
				matcher, err = domain.ReadMatcher(reader)
			}
			if err != nil {
				return
			}
			rule.DomainMatcher = matcher
			if recover {
				rule.Domain, rule.DomainSuffix = matcher.Dump()
			}
		case ruleItemDomainKeyword:
			rule.DomainKeyword, err = readRuleItemString(reader)
		case ruleItemDomainRegex:
			rule.DomainRegex, err = readRuleItemString(reader)
		case ruleItemSourceIPCIDR:
			if mmap != nil {
				rule.SourceIPSet, err = mmap.readIPSet(reader)
			} else {
				rule.SourceIPSet, err = readIPSet(reader)
			}
			if err != nil {
				return
			}
			if recover {
				rule.SourceIPCIDR = common.Map(rule.SourceIPSet.Prefixes(), func(it netip.Prefix) *badoption.Prefixable {
					return common.Ptr(badoption.Prefixable(it))
				})
			}
		case ruleItemIPCIDR:
			if mmap != nil {
				rule.IPSet, err = mmap.readIPSet(reader)
			} else {
				rule.IPSet, err = readIPSet(reader)
			}
			if err != nil {
				return
			}
			if recover {
				rule.IPCIDR = common.Map(rule.IPSet.Prefixes(), func(it netip.Prefix) *badoption.Prefixable {
					return common.Ptr(badoption.Prefixable(it))
				})
			}
		case ruleItemSourcePort:
			rule.SourcePort, err = varbin.ReadSlice[uint16](reader, binary.BigEndian)
		case ruleItemSourcePortRange:
			rule.SourcePortRange, err = readRuleItemString(reader)
		case ruleItemPort:
			rule.Port, err = varbin.ReadSlice[uint16](reader, binary.BigEndian)
		case ruleItemPortRange:
			rule.PortRange, err = readRuleItemString(reader)
		case ruleItemProcessName:
			rule.ProcessName, err = readRuleItemString(reader)
		case ruleItemProcessPath:
			rule.ProcessPath, err = readRuleItemString(reader)
		case ruleItemProcessPathRegex:
			rule.ProcessPathRegex, err = readRuleItemString(reader)
		case ruleItemPackageName:
			rule.PackageName, err = readRuleItemString(reader)
		case ruleItemPackageNameRegex:
			rule.PackageNameRegex, err = readRuleItemString(reader)
		case ruleItemWIFISSID:
			rule.WIFISSID, err = readRuleItemString(reader)
		case ruleItemWIFIBSSID:
			rule.WIFIBSSID, err = readRuleItemString(reader)
		case ruleItemAdGuardDomain:
			var matcher *domain.AdGuardMatcher
			if mmap != nil {
				matcher, err = mmap.readAdGuardMatcher(reader)
			} else {
				matcher, err = domain.ReadAdGuardMatcher(reader)
			}
			if err != nil {
				return
			}
			rule.AdGuardDomainMatcher = matcher
			if recover {
				rule.AdGuardDomain = matcher.Dump()
			}
		case ruleItemNetworkType:
			rule.NetworkType, err = varbin.ReadSlice[option.InterfaceType](reader, binary.BigEndian)
		case ruleItemNetworkIsExpensive:
			rule.NetworkIsExpensive = true
		case ruleItemNetworkIsConstrained:
			rule.NetworkIsConstrained = true
		case ruleItemNetworkInterfaceAddress:
			rule.NetworkInterfaceAddress = new(badjson.TypedMap[option.InterfaceType, badoption.Listable[*badoption.Prefixable]])
			var size uint64
			size, err = binary.ReadUvarint(reader)
			if err != nil {
				return
			}
			for range size {
				var key uint8
				err = binary.Read(reader, binary.BigEndian, &key)
				if err != nil {
					return
				}
				var value []*badoption.Prefixable
				var prefixCount uint64
				prefixCount, err = binary.ReadUvarint(reader)
				if err != nil {
					return
				}
				for j := uint64(0); j < prefixCount; j++ {
					var prefix netip.Prefix
					prefix, err = ReadPrefix(reader)
					if err != nil {
						return
					}
					value = append(value, common.Ptr(badoption.Prefixable(prefix)))
				}
				rule.NetworkInterfaceAddress.Put(option.InterfaceType(key), value)
			}
		case ruleItemDefaultInterfaceAddress:
			var value []*badoption.Prefixable
			var prefixCount uint64
			prefixCount, err = binary.ReadUvarint(reader)
			if err != nil {
				return
			}
			for j := uint64(0); j < prefixCount; j++ {
				var prefix netip.Prefix
				prefix, err = ReadPrefix(reader)
				if err != nil {
					return
				}
				value = append(value, common.Ptr(badoption.Prefixable(prefix)))
			}
			rule.DefaultInterfaceAddress = value
		case ruleItemFinal:
			err = binary.Read(reader, binary.BigEndian, &rule.Invert)
			return
		default:
			err = E.New("unknown rule item type: ", itemType, ", last type: ", lastItemType)
		}
		if err != nil {
			return
		}
		lastItemType = itemType
	}
}

func writeDefaultRule(writer varbin.Writer, rule option.DefaultHeadlessRule, generateVersion uint8, mmap *mmapWriter) error {
	err := binary.Write(writer, binary.BigEndian, uint8(0))
	if err != nil {
		return err
	}
	if len(rule.QueryType) > 0 {
		err = writeRuleItemUint16(writer, ruleItemQueryType, common.Map(rule.QueryType, func(it option.DNSQueryType) uint16 {
			return uint16(it)
		}))
		if err != nil {
			return err
		}
	}
	if len(rule.Network) > 0 {
		err = writeRuleItemString(writer, ruleItemNetwork, rule.Network)
		if err != nil {
			return err
		}
	}
	if len(rule.Domain) > 0 || len(rule.DomainSuffix) > 0 || mmap != nil && rule.DomainMatcher != nil {
		if slices.Contains(rule.Domain, "") {
			return E.New("domain: empty item is not allowed")
		}
		if slices.Contains(rule.DomainSuffix, "") {
			return E.New("domain_suffix: empty item is not allowed")
		}
		err = binary.Write(writer, binary.BigEndian, ruleItemDomain)
		if err != nil {
			return err
		}
		matcher := rule.DomainMatcher
		if len(rule.Domain) > 0 || len(rule.DomainSuffix) > 0 {
			matcher = domain.NewMatcher(rule.Domain, rule.DomainSuffix, generateVersion == C.RuleSetVersion1)
		}
		if mmap != nil {
			err = mmap.writeMatcher(writer, matcher)
		} else {
			err = matcher.Write(writer)
		}
		if err != nil {
			return err
		}
	}
	if len(rule.DomainKeyword) > 0 {
		err = writeRuleItemString(writer, ruleItemDomainKeyword, rule.DomainKeyword)
		if err != nil {
			return err
		}
	}
	if len(rule.DomainRegex) > 0 {
		err = writeRuleItemString(writer, ruleItemDomainRegex, rule.DomainRegex)
		if err != nil {
			return err
		}
	}
	if len(rule.SourceIPCIDR) > 0 || mmap != nil && rule.SourceIPSet != nil {
		err = writeRuleItemCIDR(writer, ruleItemSourceIPCIDR, rule.SourceIPCIDR, rule.SourceIPSet, mmap)
		if err != nil {
			return E.Cause(err, "source_ip_cidr")
		}
	}
	if len(rule.IPCIDR) > 0 || mmap != nil && rule.IPSet != nil {
		err = writeRuleItemCIDR(writer, ruleItemIPCIDR, rule.IPCIDR, rule.IPSet, mmap)
		if err != nil {
			return E.Cause(err, "ipcidr")
		}
	}
	if len(rule.SourcePort) > 0 {
		err = writeRuleItemUint16(writer, ruleItemSourcePort, rule.SourcePort)
		if err != nil {
			return err
		}
	}
	if len(rule.SourcePortRange) > 0 {
		err = writeRuleItemString(writer, ruleItemSourcePortRange, rule.SourcePortRange)
		if err != nil {
			return err
		}
	}
	if len(rule.Port) > 0 {
		err = writeRuleItemUint16(writer, ruleItemPort, rule.Port)
		if err != nil {
			return err
		}
	}
	if len(rule.PortRange) > 0 {
		err = writeRuleItemString(writer, ruleItemPortRange, rule.PortRange)
		if err != nil {
			return err
		}
	}
	if len(rule.ProcessName) > 0 {
		err = writeRuleItemString(writer, ruleItemProcessName, rule.ProcessName)
		if err != nil {
			return err
		}
	}
	if len(rule.ProcessPath) > 0 {
		err = writeRuleItemString(writer, ruleItemProcessPath, rule.ProcessPath)
		if err != nil {
			return err
		}
	}
	if len(rule.ProcessPathRegex) > 0 {
		err = writeRuleItemString(writer, ruleItemProcessPathRegex, rule.ProcessPathRegex)
		if err != nil {
			return err
		}
	}
	if len(rule.PackageName) > 0 {
		err = writeRuleItemString(writer, ruleItemPackageName, rule.PackageName)
		if err != nil {
			return err
		}
	}
	if len(rule.PackageNameRegex) > 0 {
		if generateVersion < C.RuleSetVersion5 {
			return E.New("`package_name_regex` rule item is only supported in version 5 or later")
		}
		err = writeRuleItemString(writer, ruleItemPackageNameRegex, rule.PackageNameRegex)
		if err != nil {
			return err
		}
	}
	if len(rule.NetworkType) > 0 {
		if generateVersion < C.RuleSetVersion3 {
			return E.New("`network_type` rule item is only supported in version 3 or later")
		}
		err = writeRuleItemUint8(writer, ruleItemNetworkType, rule.NetworkType)
		if err != nil {
			return err
		}
	}
	if rule.NetworkIsExpensive {
		if generateVersion < C.RuleSetVersion3 {
			return E.New("`network_is_expensive` rule item is only supported in version 3 or later")
		}
		err = binary.Write(writer, binary.BigEndian, ruleItemNetworkIsExpensive)
		if err != nil {
			return err
		}
	}
	if rule.NetworkIsConstrained {
		if generateVersion < C.RuleSetVersion3 {
			return E.New("`network_is_constrained` rule item is only supported in version 3 or later")
		}
		err = binary.Write(writer, binary.BigEndian, ruleItemNetworkIsConstrained)
		if err != nil {
			return err
		}
	}
	if rule.NetworkInterfaceAddress != nil && rule.NetworkInterfaceAddress.Size() > 0 {
		if generateVersion < C.RuleSetVersion4 {
			return E.New("`network_interface_address` rule item is only supported in version 4 or later")
		}
		err = writer.WriteByte(ruleItemNetworkInterfaceAddress)
		if err != nil {
			return err
		}
		_, err = varbin.WriteUvarint(writer, uint64(rule.NetworkInterfaceAddress.Size()))
		if err != nil {
			return err
		}
		for _, entry := range rule.NetworkInterfaceAddress.Entries() {
			err = binary.Write(writer, binary.BigEndian, uint8(entry.Key.Build()))
			if err != nil {
				return err
			}
			_, err = varbin.WriteUvarint(writer, uint64(len(entry.Value)))
			if err != nil {
				return err
			}
			for _, rawPrefix := range entry.Value {
				err = WritePrefix(writer, rawPrefix.Build(netip.Prefix{}))
				if err != nil {
					return err
				}
			}
		}
	}
	if len(rule.DefaultInterfaceAddress) > 0 {
		if generateVersion < C.RuleSetVersion4 {
			return E.New("`default_interface_address` rule item is only supported in version 4 or later")
		}
		err = writer.WriteByte(ruleItemDefaultInterfaceAddress)
		if err != nil {
			return err
		}
		_, err = varbin.WriteUvarint(writer, uint64(len(rule.DefaultInterfaceAddress)))
		if err != nil {
			return err
		}
		for _, rawPrefix := range rule.DefaultInterfaceAddress {
			err = WritePrefix(writer, rawPrefix.Build(netip.Prefix{}))
			if err != nil {
				return err
			}
		}
	}
	if len(rule.WIFISSID) > 0 {
		err = writeRuleItemString(writer, ruleItemWIFISSID, rule.WIFISSID)
		if err != nil {
			return err
		}
	}
	if len(rule.WIFIBSSID) > 0 {
		err = writeRuleItemString(writer, ruleItemWIFIBSSID, rule.WIFIBSSID)
		if err != nil {
			return err
		}
	}
	if len(rule.AdGuardDomain) > 0 || mmap != nil && rule.AdGuardDomainMatcher != nil {
		if generateVersion < C.RuleSetVersion2 {
			return E.New("AdGuard rule items is only supported in version 2 or later")
		}
		err = binary.Write(writer, binary.BigEndian, ruleItemAdGuardDomain)
		if err != nil {
			return err
		}
		matcher := rule.AdGuardDomainMatcher
		if len(rule.AdGuardDomain) > 0 {
			matcher = domain.NewAdGuardMatcher(rule.AdGuardDomain)
		}
		if mmap != nil {
			err = mmap.writeAdGuardMatcher(writer, matcher)
		} else {
			err = matcher.Write(writer)
		}
		if err != nil {
			return err
		}
	}
	err = binary.Write(writer, binary.BigEndian, ruleItemFinal)
	if err != nil {
		return err
	}
	err = binary.Write(writer, binary.BigEndian, rule.Invert)
	if err != nil {
		return err
	}
	return nil
}

func readRuleItemString(reader varbin.Reader) ([]string, error) {
	length, err := binary.ReadUvarint(reader)
	if err != nil {
		return nil, err
	}
	var result []string
	for range length {
		var value []byte
		value, err = varbin.ReadSlice[byte](reader, binary.BigEndian)
		if err != nil {
			return nil, err
		}
		result = append(result, string(value))
	}
	return result, nil
}

func writeRuleItemString(writer varbin.Writer, itemType uint8, value []string) error {
	err := writer.WriteByte(itemType)
	if err != nil {
		return err
	}
	_, err = varbin.WriteUvarint(writer, uint64(len(value)))
	if err != nil {
		return err
	}
	for _, s := range value {
		_, err = varbin.WriteUvarint(writer, uint64(len(s)))
		if err != nil {
			return err
		}
		_, err = writer.Write([]byte(s))
		if err != nil {
			return err
		}
	}
	return nil
}

func writeRuleItemUint8[E ~uint8](writer varbin.Writer, itemType uint8, value []E) error {
	err := writer.WriteByte(itemType)
	if err != nil {
		return err
	}
	_, err = varbin.WriteUvarint(writer, uint64(len(value)))
	if err != nil {
		return err
	}
	_, err = writer.Write(*(*[]byte)(unsafe.Pointer(&value)))
	return err
}

func writeRuleItemUint16(writer varbin.Writer, itemType uint8, value []uint16) error {
	err := writer.WriteByte(itemType)
	if err != nil {
		return err
	}
	_, err = varbin.WriteUvarint(writer, uint64(len(value)))
	if err != nil {
		return err
	}
	return binary.Write(writer, binary.BigEndian, value)
}

func writeRuleItemCIDR(writer varbin.Writer, itemType uint8, value []*badoption.Prefixable, rawSet *ipset.Set, mmap *mmapWriter) error {
	set := rawSet
	if len(value) > 0 {
		var builder netipx.IPSetBuilder
		for _, prefixable := range value {
			builder.AddPrefix(prefixable.Build(netip.Prefix{}))
		}
		ipSet, err := builder.IPSet()
		if err != nil {
			return err
		}
		set = ipset.FromIPSet(ipSet)
	}
	err := binary.Write(writer, binary.BigEndian, itemType)
	if err != nil {
		return err
	}
	if mmap != nil {
		return mmap.writeIPSet(writer, set)
	}
	return writeIPSet(writer, set)
}

func readLogicalRule(reader varbin.Reader, recovery bool, depth int, mmap *mmapReader) (logicalRule option.LogicalHeadlessRule, err error) {
	mode, err := reader.ReadByte()
	if err != nil {
		return
	}
	switch mode {
	case 0:
		logicalRule.Mode = C.LogicalTypeAnd
	case 1:
		logicalRule.Mode = C.LogicalTypeOr
	default:
		err = E.New("unknown logical mode: ", mode)
		return
	}
	length, err := binary.ReadUvarint(reader)
	if err != nil {
		return
	}
	for i := range length {
		var rule option.HeadlessRule
		rule, err = readRule(reader, recovery, depth+1, mmap)
		if err != nil {
			err = E.Cause(err, "read logical rule [", i, "]")
			return
		}
		logicalRule.Rules = append(logicalRule.Rules, rule)
	}
	err = binary.Read(reader, binary.BigEndian, &logicalRule.Invert)
	if err != nil {
		return
	}
	return
}

func writeLogicalRule(writer varbin.Writer, logicalRule option.LogicalHeadlessRule, generateVersion uint8, mmap *mmapWriter) error {
	err := binary.Write(writer, binary.BigEndian, uint8(1))
	if err != nil {
		return err
	}
	switch logicalRule.Mode {
	case C.LogicalTypeAnd:
		err = binary.Write(writer, binary.BigEndian, uint8(0))
	case C.LogicalTypeOr:
		err = binary.Write(writer, binary.BigEndian, uint8(1))
	default:
		return E.New("unknown logical mode: ", logicalRule.Mode)
	}
	if err != nil {
		return err
	}
	_, err = varbin.WriteUvarint(writer, uint64(len(logicalRule.Rules)))
	if err != nil {
		return err
	}
	for _, rule := range logicalRule.Rules {
		err = writeRule(writer, rule, generateVersion, mmap)
		if err != nil {
			return err
		}
	}
	err = binary.Write(writer, binary.BigEndian, logicalRule.Invert)
	if err != nil {
		return err
	}
	return nil
}