github.com/imannamdari/v2ray-core/v5@v5.0.5/infra/conf/geodata/memconservative/cache.go (about)

     1  package memconservative
     2  
     3  import (
     4  	"io/ioutil"
     5  	"strings"
     6  
     7  	"google.golang.org/protobuf/proto"
     8  
     9  	"github.com/imannamdari/v2ray-core/v5/app/router/routercommon"
    10  	"github.com/imannamdari/v2ray-core/v5/common/platform"
    11  )
    12  
    13  type GeoIPCache map[string]*routercommon.GeoIP
    14  
    15  func (g GeoIPCache) Has(key string) bool {
    16  	return !(g.Get(key) == nil)
    17  }
    18  
    19  func (g GeoIPCache) Get(key string) *routercommon.GeoIP {
    20  	if g == nil {
    21  		return nil
    22  	}
    23  	return g[key]
    24  }
    25  
    26  func (g GeoIPCache) Set(key string, value *routercommon.GeoIP) {
    27  	if g == nil {
    28  		g = make(map[string]*routercommon.GeoIP)
    29  	}
    30  	g[key] = value
    31  }
    32  
    33  func (g GeoIPCache) Unmarshal(filename, code string) (*routercommon.GeoIP, error) {
    34  	asset := platform.GetAssetLocation(filename)
    35  	idx := strings.ToLower(asset + ":" + code)
    36  	if g.Has(idx) {
    37  		return g.Get(idx), nil
    38  	}
    39  
    40  	geoipBytes, err := Decode(asset, code)
    41  	switch err {
    42  	case nil:
    43  		var geoip routercommon.GeoIP
    44  		if err := proto.Unmarshal(geoipBytes, &geoip); err != nil {
    45  			return nil, err
    46  		}
    47  		g.Set(idx, &geoip)
    48  		return &geoip, nil
    49  
    50  	case errCodeNotFound:
    51  		return nil, newError("country code ", code, " not found in ", filename)
    52  
    53  	case errFailedToReadBytes, errFailedToReadExpectedLenBytes,
    54  		errInvalidGeodataFile, errInvalidGeodataVarintLength:
    55  		newError("failed to decode geoip file: ", filename, ", fallback to the original ReadFile method")
    56  		geoipBytes, err = ioutil.ReadFile(asset)
    57  		if err != nil {
    58  			return nil, err
    59  		}
    60  		var geoipList routercommon.GeoIPList
    61  		if err := proto.Unmarshal(geoipBytes, &geoipList); err != nil {
    62  			return nil, err
    63  		}
    64  		for _, geoip := range geoipList.GetEntry() {
    65  			if strings.EqualFold(code, geoip.GetCountryCode()) {
    66  				g.Set(idx, geoip)
    67  				return geoip, nil
    68  			}
    69  		}
    70  
    71  	default:
    72  		return nil, err
    73  	}
    74  
    75  	return nil, newError("country code ", code, " not found in ", filename)
    76  }
    77  
    78  type GeoSiteCache map[string]*routercommon.GeoSite
    79  
    80  func (g GeoSiteCache) Has(key string) bool {
    81  	return !(g.Get(key) == nil)
    82  }
    83  
    84  func (g GeoSiteCache) Get(key string) *routercommon.GeoSite {
    85  	if g == nil {
    86  		return nil
    87  	}
    88  	return g[key]
    89  }
    90  
    91  func (g GeoSiteCache) Set(key string, value *routercommon.GeoSite) {
    92  	if g == nil {
    93  		g = make(map[string]*routercommon.GeoSite)
    94  	}
    95  	g[key] = value
    96  }
    97  
    98  func (g GeoSiteCache) Unmarshal(filename, code string) (*routercommon.GeoSite, error) {
    99  	asset := platform.GetAssetLocation(filename)
   100  	idx := strings.ToLower(asset + ":" + code)
   101  	if g.Has(idx) {
   102  		return g.Get(idx), nil
   103  	}
   104  
   105  	geositeBytes, err := Decode(asset, code)
   106  	switch err {
   107  	case nil:
   108  		var geosite routercommon.GeoSite
   109  		if err := proto.Unmarshal(geositeBytes, &geosite); err != nil {
   110  			return nil, err
   111  		}
   112  		g.Set(idx, &geosite)
   113  		return &geosite, nil
   114  
   115  	case errCodeNotFound:
   116  		return nil, newError("list ", code, " not found in ", filename)
   117  
   118  	case errFailedToReadBytes, errFailedToReadExpectedLenBytes,
   119  		errInvalidGeodataFile, errInvalidGeodataVarintLength:
   120  		newError("failed to decode geoip file: ", filename, ", fallback to the original ReadFile method")
   121  		geositeBytes, err = ioutil.ReadFile(asset)
   122  		if err != nil {
   123  			return nil, err
   124  		}
   125  		var geositeList routercommon.GeoSiteList
   126  		if err := proto.Unmarshal(geositeBytes, &geositeList); err != nil {
   127  			return nil, err
   128  		}
   129  		for _, geosite := range geositeList.GetEntry() {
   130  			if strings.EqualFold(code, geosite.GetCountryCode()) {
   131  				g.Set(idx, geosite)
   132  				return geosite, nil
   133  			}
   134  		}
   135  
   136  	default:
   137  		return nil, err
   138  	}
   139  
   140  	return nil, newError("list ", code, " not found in ", filename)
   141  }