github.com/metacubex/mihomo@v1.18.5/component/geodata/memconservative/cache.go (about)

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