memconservative geofile loader

infra/conf/geodata/memconservative/cache.go

@@ -0,0 +1,139 @@
+package memconservative
+
+import (
+	"github.com/golang/protobuf/proto"
+	"github.com/v2fly/v2ray-core/v4/app/router"
+	"github.com/v2fly/v2ray-core/v4/common/platform"
+	"io/ioutil"
+	"strings"
+)
+
+type GeoIPCache map[string]*router.GeoIP
+
+func (g GeoIPCache) Has(key string) bool {
+	return !(g.Get(key) == nil)
+}
+
+func (g GeoIPCache) Get(key string) *router.GeoIP {
+	if g == nil {
+		return nil
+	}
+	return g[key]
+}
+
+func (g GeoIPCache) Set(key string, value *router.GeoIP) {
+	if g == nil {
+		g = make(map[string]*router.GeoIP)
+	}
+	g[key] = value
+}
+
+func (g GeoIPCache) Unmarshal(filename, code string) (*router.GeoIP, error) {
+	asset := platform.GetAssetLocation(filename)
+	idx := strings.ToUpper(asset + "|" + code)
+	if g.Has(idx) {
+		return g.Get(idx), nil
+	}
+
+	geoipBytes, err := Decode(asset, code)
+	switch err {
+	case nil:
+		var geoip router.GeoIP
+		if err := proto.Unmarshal(geoipBytes, &geoip); err != nil {
+			return nil, err
+		}
+		g.Set(idx, &geoip)
+		return &geoip, nil
+
+	case errCodeNotFound:
+		return nil, newError(code, " not found in ", filename)
+
+	case errFailedToReadBytes, errFailedToReadExpectedLenBytes,
+		errInvalidGeodataFile, errInvalidGeodataVarintLength:
+		newError("failed to decode geodata file: ", filename, ". Fallback to the original ReadFile method.").AtWarning().WriteToLog()
+		geoipBytes, err = ioutil.ReadFile(asset)
+		if err != nil {
+			return nil, err
+		}
+		var geoipList router.GeoIPList
+		if err := proto.Unmarshal(geoipBytes, &geoipList); err != nil {
+			return nil, err
+		}
+		for _, geoip := range geoipList.GetEntry() {
+			if strings.EqualFold(code, geoip.GetCountryCode()) {
+				g.Set(idx, geoip)
+				return geoip, nil
+			}
+		}
+
+	default:
+		return nil, err
+	}
+
+	return nil, newError(code, " not found in ", filename)
+}
+
+type GeoSiteCache map[string]*router.GeoSite
+
+func (g GeoSiteCache) Has(key string) bool {
+	return !(g.Get(key) == nil)
+}
+
+func (g GeoSiteCache) Get(key string) *router.GeoSite {
+	if g == nil {
+		return nil
+	}
+	return g[key]
+}
+
+func (g GeoSiteCache) Set(key string, value *router.GeoSite) {
+	if g == nil {
+		g = make(map[string]*router.GeoSite)
+	}
+	g[key] = value
+}
+
+func (g GeoSiteCache) Unmarshal(filename, code string) (*router.GeoSite, error) {
+	asset := platform.GetAssetLocation(filename)
+	idx := strings.ToUpper(asset + "|" + code)
+	if g.Has(idx) {
+		return g.Get(idx), nil
+	}
+
+	geositeBytes, err := Decode(asset, code)
+	switch err {
+	case nil:
+		var geosite router.GeoSite
+		if err := proto.Unmarshal(geositeBytes, &geosite); err != nil {
+			return nil, err
+		}
+		g.Set(idx, &geosite)
+		return &geosite, nil
+
+	case errCodeNotFound:
+		return nil, newError(code, " not found in ", filename)
+
+	case errFailedToReadBytes, errFailedToReadExpectedLenBytes,
+		errInvalidGeodataFile, errInvalidGeodataVarintLength:
+		newError("failed to decode geodata file: ", filename, ". Fallback to the original ReadFile method.").AtWarning().WriteToLog()
+		geositeBytes, err = ioutil.ReadFile(asset)
+		if err != nil {
+			return nil, err
+		}
+		var geositeList router.GeoSiteList
+		if err := proto.Unmarshal(geositeBytes, &geositeList); err != nil {
+			return nil, err
+		}
+		for _, geosite := range geositeList.GetEntry() {
+			if strings.EqualFold(code, geosite.GetCountryCode()) {
+				g.Set(idx, geosite)
+				return geosite, nil
+			}
+		}
+
+	default:
+		return nil, err
+	}
+
+	return nil, newError(code, " not found in ", filename)
+}

infra/conf/geodata/memconservative/decode.go

@@ -0,0 +1,104 @@
+package memconservative
+
+import (
+	"errors"
+	"github.com/v2fly/v2ray-core/v4/common/platform/filesystem"
+	"google.golang.org/protobuf/encoding/protowire"
+	"io"
+	"strings"
+)
+
+var (
+	errFailedToReadBytes            = errors.New("failed to read bytes")
+	errFailedToReadExpectedLenBytes = errors.New("failed to read expected length of bytes")
+	errInvalidGeodataFile           = errors.New("invalid geodata file")
+	errInvalidGeodataVarintLength   = errors.New("invalid geodata varint length")
+	errCodeNotFound                 = errors.New("code not found")
+)
+
+func emitBytes(f io.ReadSeeker, code string) ([]byte, error) {
+	count := 1
+	isInner := false
+	tempContainer := make([]byte, 0, 5)
+
+	var result []byte
+	var advancedN uint64 = 1
+	var geoDataVarintLength, codeVarintLength, varintLenByteLen uint64 = 0, 0, 0
+
+Loop:
+	for {
+		container := make([]byte, advancedN)
+		bytesRead, err := f.Read(container)
+		if err == io.EOF {
+			return nil, errCodeNotFound
+		}
+		if err != nil {
+			return nil, errFailedToReadBytes
+		}
+		if bytesRead != len(container) {
+			return nil, errFailedToReadExpectedLenBytes
+		}
+
+		switch count {
+		case 1, 3: // data type ((field_number << 3) | wire_type)
+			if container[0] != 10 { // byte `0A` equals to `10` in decimal
+				return nil, errInvalidGeodataFile
+			}
+			advancedN = 1
+			count++
+		case 2, 4: // data length
+			tempContainer = append(tempContainer, container...)
+			if container[0] > 127 { // max one-byte-length byte `7F`(0FFF FFFF) equals to `127` in decimal
+				advancedN = 1
+				goto Loop
+			}
+			lenVarint, n := protowire.ConsumeVarint(tempContainer)
+			if n < 0 {
+				return nil, errInvalidGeodataVarintLength
+			}
+			tempContainer = nil
+			if !isInner {
+				isInner = true
+				geoDataVarintLength = lenVarint
+				advancedN = 1
+			} else {
+				isInner = false
+				codeVarintLength = lenVarint
+				varintLenByteLen = uint64(n)
+				advancedN = codeVarintLength
+			}
+			count++
+		case 5: // data value
+			if strings.EqualFold(string(container), code) {
+				count++
+				offset := -(1 + int64(varintLenByteLen) + int64(codeVarintLength))
+				f.Seek(offset, 1)               // back to the start of GeoIP or GeoSite varint
+				advancedN = geoDataVarintLength // the number of bytes to be read in next round
+			} else {
+				count = 1
+				offset := int64(geoDataVarintLength) - int64(codeVarintLength) - int64(varintLenByteLen) - 1
+				f.Seek(offset, 1) // skip the unmatched GeoIP or GeoSite varint
+				advancedN = 1     // the next round will be the start of another GeoIPList or GeoSiteList
+			}
+		case 6: // matched GeoIP or GeoSite varint
+			result = container
+			break Loop
+		}
+
+	}
+	return result, nil
+}
+
+func Decode(filename, code string) ([]byte, error) {
+	f, err := filesystem.NewFileSeeker(filename)
+	if err != nil {
+		return nil, newError("failed to open file: ", filename).Base(err)
+	}
+	defer f.Close()
+
+	geoBytes, err := emitBytes(f, code)
+	if err != nil {
+		return nil, err
+	}
+	return geoBytes, nil
+}

infra/conf/geodata/memconservative/decode_test.go

@@ -0,0 +1,73 @@
+package memconservative
+
+import (
+	"errors"
+	"github.com/google/go-cmp/cmp"
+	"github.com/v2fly/v2ray-core/v4/common"
+	"github.com/v2fly/v2ray-core/v4/common/platform"
+	"github.com/v2fly/v2ray-core/v4/common/platform/filesystem"
+	"io/fs"
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+const (
+	geoipURL   = "https://raw.githubusercontent.com/v2fly/geoip/release/geoip.dat"
+	geositeURL = "https://raw.githubusercontent.com/v2fly/domain-list-community/release/dlc.dat"
+)
+
+func init() {
+	wd, err := os.Getwd()
+	common.Must(err)
+
+	tempPath := filepath.Join(wd, "..", "..", "testing", "temp")
+	geoipPath := filepath.Join(tempPath, "geoip.dat")
+	geositePath := filepath.Join(tempPath, "geosite.dat")
+
+	os.Setenv("v2ray.location.asset", tempPath)
+
+	common.Must(os.MkdirAll(tempPath, 0755))
+
+	if _, err := os.Stat(platform.GetAssetLocation("geoip.dat")); err != nil && errors.Is(err, fs.ErrNotExist) {
+		if _, err := os.Stat(geoipPath); err != nil && errors.Is(err, fs.ErrNotExist) {
+			geoipBytes, err := common.FetchHTTPContent(geoipURL)
+			common.Must(err)
+			common.Must(filesystem.WriteFile(geoipPath, geoipBytes))
+		}
+	}
+
+	if _, err := os.Stat(platform.GetAssetLocation("geosite.dat")); err != nil && errors.Is(err, fs.ErrNotExist) {
+		if _, err := os.Stat(geositePath); err != nil && errors.Is(err, fs.ErrNotExist) {
+			geositeBytes, err := common.FetchHTTPContent(geositeURL)
+			common.Must(err)
+			common.Must(filesystem.WriteFile(geositePath, geositeBytes))
+		}
+	}
+}
+
+func TestDecodeGeoIP(t *testing.T) {
+	filename := platform.GetAssetLocation("geoip.dat")
+	result, err := Decode(filename, "test")
+	if err != nil {
+		t.Error(err)
+	}
+
+	expected := []byte{10, 4, 84, 69, 83, 84, 18, 8, 10, 4, 127, 0, 0, 0, 16, 8}
+	if cmp.Diff(result, expected) != "" {
+		t.Errorf("failed to load geoip:test, expected: %v, got: %v", expected, result)
+	}
+}
+
+func TestDecodeGeoSite(t *testing.T) {
+	filename := platform.GetAssetLocation("geosite.dat")
+	result, err := Decode(filename, "test")
+	if err != nil {
+		t.Error(err)
+	}
+
+	expected := []byte{10, 4, 84, 69, 83, 84, 18, 20, 8, 3, 18, 16, 116, 101, 115, 116, 46, 101, 120, 97, 109, 112, 108, 101, 46, 99, 111, 109}
+	if cmp.Diff(result, expected) != "" {
+		t.Errorf("failed to load geosite:test, expected: %v, got: %v", expected, result)
+	}
+}

infra/conf/geodata/memconservative/errors.generated.go

@@ -0,0 +1,9 @@
+package memconservative
+
+import "github.com/v2fly/v2ray-core/v4/common/errors"
+
+type errPathObjHolder struct{}
+
+func newError(values ...interface{}) *errors.Error {
+	return errors.New(values...).WithPathObj(errPathObjHolder{})
+}

infra/conf/geodata/memconservative/memc.go

@@ -0,0 +1,42 @@
+package memconservative
+
+import (
+	"github.com/v2fly/v2ray-core/v4/app/router"
+	"github.com/v2fly/v2ray-core/v4/infra/conf/geodata"
+	"runtime"
+)
+
+//go:generate go run github.com/v2fly/v2ray-core/v4/common/errors/errorgen
+
+type memConservativeLoader struct {
+	geoipcache   GeoIPCache
+	geositecache GeoSiteCache
+}
+
+func (m *memConservativeLoader) LoadIP(filename, country string) ([]*router.CIDR, error) {
+	defer runtime.GC()
+	geoip, err := m.geoipcache.Unmarshal(filename, country)
+	if err != nil {
+		return nil, newError("failed to decode geodata file: ", filename).Base(err)
+	}
+	return geoip.Cidr, nil
+}
+
+func (m *memConservativeLoader) LoadSite(filename, list string) ([]*router.Domain, error) {
+	defer runtime.GC()
+	geosite, err := m.geositecache.Unmarshal(filename, list)
+	if err != nil {
+		return nil, newError("failed to decode geodata file: ", filename).Base(err)
+	}
+	return geosite.Domain, nil
+}
+
+func newMemConservativeLoader() geodata.LoaderImplementation {
+	return &memConservativeLoader{make(map[string]*router.GeoIP), make(map[string]*router.GeoSite)}
+}
+
+func init() {
+	geodata.RegisterGeoDataLoaderImplementationCreator("memconservative", func() geodata.LoaderImplementation {
+		return newMemConservativeLoader()
+	})
+}