Browse Source

Add geosite protocol

世界 3 years ago
parent
commit
f76102dab5
3 changed files with 223 additions and 0 deletions
  1. 97 0
      common/geosite/reader.go
  2. 62 0
      common/geosite/rule.go
  3. 64 0
      common/geosite/writer.go

+ 97 - 0
common/geosite/reader.go

@@ -0,0 +1,97 @@
+package geosite
+
+import (
+	"io"
+	"sync"
+
+	E "github.com/sagernet/sing/common/exceptions"
+	"github.com/sagernet/sing/common/rw"
+)
+
+type Reader struct {
+	reader       io.ReadSeeker
+	access       sync.Mutex
+	metadataRead bool
+	domainIndex  map[string]int
+	domainLength map[string]int
+}
+
+func (r *Reader) readMetadata() error {
+	version, err := rw.ReadByte(r.reader)
+	if err != nil {
+		return err
+	}
+	if version != 0 {
+		return E.New("unknown version")
+	}
+	entryLength, err := rw.ReadUVariant(r.reader)
+	if err != nil {
+		return err
+	}
+	keys := make([]string, entryLength)
+	domainIndex := make(map[string]int)
+	domainLength := make(map[string]int)
+	for i := 0; i < int(entryLength); i++ {
+		var (
+			code       string
+			codeIndex  uint64
+			codeLength uint64
+		)
+		code, err = rw.ReadVString(r.reader)
+		if err != nil {
+			return err
+		}
+		keys[i] = code
+		codeIndex, err = rw.ReadUVariant(r.reader)
+		if err != nil {
+			return err
+		}
+		codeLength, err = rw.ReadUVariant(r.reader)
+		if err != nil {
+			return err
+		}
+		domainIndex[code] = int(codeIndex)
+		domainLength[code] = int(codeLength)
+	}
+	r.domainIndex = domainIndex
+	r.domainLength = domainLength
+	r.metadataRead = true
+	return nil
+}
+
+func (r *Reader) Read(code string) ([]Item, error) {
+	r.access.Lock()
+	defer r.access.Unlock()
+	if !r.metadataRead {
+		err := r.readMetadata()
+		if err != nil {
+			return nil, err
+		}
+	}
+	if _, exists := r.domainIndex[code]; !exists {
+		return nil, E.New("code ", code, " not exists!")
+	}
+	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
+	}
+	_, err := r.reader.Seek(int64(r.domainIndex[code])-counter.Count(), io.SeekCurrent)
+	return domain, err
+}
+
+func (r *Reader) Upstream() any {
+	return r.reader
+}

+ 62 - 0
common/geosite/rule.go

@@ -0,0 +1,62 @@
+package geosite
+
+import "github.com/sagernet/sing-box/option"
+
+type ItemType = uint8
+
+const (
+	RuleTypeDomain ItemType = iota
+	RuleTypeDomainSuffix
+	RuleTypeDomainKeyword
+	RuleTypeDomainRegex
+)
+
+type Item struct {
+	Type  ItemType
+	Value string
+}
+
+func Compile(code []Item) option.DefaultRule {
+	var domainLength int
+	var domainSuffixLength int
+	var domainKeywordLength int
+	var domainRegexLength int
+	for _, item := range code {
+		switch item.Type {
+		case RuleTypeDomain:
+			domainLength++
+		case RuleTypeDomainSuffix:
+			domainSuffixLength++
+		case RuleTypeDomainKeyword:
+			domainKeywordLength++
+		case RuleTypeDomainRegex:
+			domainRegexLength++
+		}
+	}
+	var codeRule option.DefaultRule
+	if domainLength > 0 {
+		codeRule.Domain = make([]string, 0, domainLength)
+	}
+	if domainSuffixLength > 0 {
+		codeRule.DomainSuffix = make([]string, 0, domainSuffixLength)
+	}
+	if domainKeywordLength > 0 {
+		codeRule.DomainKeyword = make([]string, 0, domainKeywordLength)
+	}
+	if domainRegexLength > 0 {
+		codeRule.DomainRegex = make([]string, 0, domainRegexLength)
+	}
+	for _, item := range code {
+		switch item.Type {
+		case RuleTypeDomain:
+			codeRule.Domain = append(codeRule.Domain, item.Value)
+		case RuleTypeDomainSuffix:
+			codeRule.DomainSuffix = append(codeRule.DomainSuffix, item.Value)
+		case RuleTypeDomainKeyword:
+			codeRule.DomainKeyword = append(codeRule.DomainKeyword, item.Value)
+		case RuleTypeDomainRegex:
+			codeRule.DomainRegex = append(codeRule.DomainRegex, item.Value)
+		}
+	}
+	return codeRule
+}

+ 64 - 0
common/geosite/writer.go

@@ -0,0 +1,64 @@
+package geosite
+
+import (
+	"bytes"
+	"io"
+	"sort"
+
+	"github.com/sagernet/sing/common/rw"
+)
+
+func Write(writer io.Writer, domains map[string][]Item) error {
+	keys := make([]string, 0, len(domains))
+	for code := range domains {
+		keys = append(keys, code)
+	}
+	sort.Strings(keys)
+
+	content := &bytes.Buffer{}
+	index := make(map[string]int)
+	for _, code := range keys {
+		index[code] = content.Len()
+		for _, domain := range domains[code] {
+			err := rw.WriteByte(content, byte(domain.Type))
+			if err != nil {
+				return err
+			}
+			if err = rw.WriteVString(content, domain.Value); err != nil {
+				return err
+			}
+		}
+	}
+
+	err := rw.WriteByte(writer, 0)
+	if err != nil {
+		return err
+	}
+
+	err = rw.WriteUVariant(writer, uint64(len(keys)))
+	if err != nil {
+		return err
+	}
+
+	for _, code := range keys {
+		err = rw.WriteVString(writer, code)
+		if err != nil {
+			return err
+		}
+		err = rw.WriteUVariant(writer, uint64(index[code]))
+		if err != nil {
+			return err
+		}
+		err = rw.WriteUVariant(writer, uint64(len(domains[code])))
+		if err != nil {
+			return err
+		}
+	}
+
+	_, err = writer.Write(content.Bytes())
+	if err != nil {
+		return err
+	}
+
+	return nil
+}