mirror of https://github.com/v2fly/v2ray-core.git
memconservative geofile loader
This commit is contained in:
Shelikhoo
parent
0f1fac84ca
commit
4610c5e23f
|
|
@ -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)
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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{})
|
||||
}
|
||||
|
|
@ -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()
|
||||
})
|
||||
}
|
||||
Loading…
Reference in New Issue