github.com/yaling888/clash@v1.53.0/component/geodata/memconservative/cache.go (about)

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