github.com/metacubex/mihomo@v1.18.5/component/mmdb/mmdb.go (about)

     1  package mmdb
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"net/http"
     7  	"os"
     8  	"sync"
     9  	"time"
    10  
    11  	mihomoOnce "github.com/metacubex/mihomo/common/once"
    12  	mihomoHttp "github.com/metacubex/mihomo/component/http"
    13  	C "github.com/metacubex/mihomo/constant"
    14  	"github.com/metacubex/mihomo/log"
    15  
    16  	"github.com/oschwald/maxminddb-golang"
    17  )
    18  
    19  type databaseType = uint8
    20  
    21  const (
    22  	typeMaxmind databaseType = iota
    23  	typeSing
    24  	typeMetaV0
    25  )
    26  
    27  var (
    28  	IPreader  IPReader
    29  	ASNreader ASNReader
    30  	IPonce    sync.Once
    31  	ASNonce   sync.Once
    32  )
    33  
    34  func LoadFromBytes(buffer []byte) {
    35  	IPonce.Do(func() {
    36  		mmdb, err := maxminddb.FromBytes(buffer)
    37  		if err != nil {
    38  			log.Fatalln("Can't load mmdb: %s", err.Error())
    39  		}
    40  		IPreader = IPReader{Reader: mmdb}
    41  		switch mmdb.Metadata.DatabaseType {
    42  		case "sing-geoip":
    43  			IPreader.databaseType = typeSing
    44  		case "Meta-geoip0":
    45  			IPreader.databaseType = typeMetaV0
    46  		default:
    47  			IPreader.databaseType = typeMaxmind
    48  		}
    49  	})
    50  }
    51  
    52  func Verify(path string) bool {
    53  	instance, err := maxminddb.Open(path)
    54  	if err == nil {
    55  		instance.Close()
    56  	}
    57  	return err == nil
    58  }
    59  
    60  func IPInstance() IPReader {
    61  	IPonce.Do(func() {
    62  		mmdbPath := C.Path.MMDB()
    63  		log.Infoln("Load MMDB file: %s", mmdbPath)
    64  		mmdb, err := maxminddb.Open(mmdbPath)
    65  		if err != nil {
    66  			log.Fatalln("Can't load MMDB: %s", err.Error())
    67  		}
    68  		IPreader = IPReader{Reader: mmdb}
    69  		switch mmdb.Metadata.DatabaseType {
    70  		case "sing-geoip":
    71  			IPreader.databaseType = typeSing
    72  		case "Meta-geoip0":
    73  			IPreader.databaseType = typeMetaV0
    74  		default:
    75  			IPreader.databaseType = typeMaxmind
    76  		}
    77  	})
    78  
    79  	return IPreader
    80  }
    81  
    82  func DownloadMMDB(path string) (err error) {
    83  	ctx, cancel := context.WithTimeout(context.Background(), time.Second*90)
    84  	defer cancel()
    85  	resp, err := mihomoHttp.HttpRequest(ctx, C.MmdbUrl, http.MethodGet, http.Header{"User-Agent": {C.UA}}, nil)
    86  	if err != nil {
    87  		return
    88  	}
    89  	defer resp.Body.Close()
    90  
    91  	f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0o644)
    92  	if err != nil {
    93  		return err
    94  	}
    95  	defer f.Close()
    96  	_, err = io.Copy(f, resp.Body)
    97  
    98  	return err
    99  }
   100  
   101  func ASNInstance() ASNReader {
   102  	ASNonce.Do(func() {
   103  		ASNPath := C.Path.ASN()
   104  		log.Infoln("Load ASN file: %s", ASNPath)
   105  		asn, err := maxminddb.Open(ASNPath)
   106  		if err != nil {
   107  			log.Fatalln("Can't load ASN: %s", err.Error())
   108  		}
   109  		ASNreader = ASNReader{Reader: asn}
   110  	})
   111  
   112  	return ASNreader
   113  }
   114  
   115  func DownloadASN(path string) (err error) {
   116  	ctx, cancel := context.WithTimeout(context.Background(), time.Second*90)
   117  	defer cancel()
   118  	resp, err := mihomoHttp.HttpRequest(ctx, C.ASNUrl, http.MethodGet, http.Header{"User-Agent": {C.UA}}, nil)
   119  	if err != nil {
   120  		return
   121  	}
   122  	defer resp.Body.Close()
   123  
   124  	f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0o644)
   125  	if err != nil {
   126  		return err
   127  	}
   128  	defer f.Close()
   129  	_, err = io.Copy(f, resp.Body)
   130  
   131  	return err
   132  }
   133  
   134  func ReloadIP() {
   135  	mihomoOnce.Reset(&IPonce)
   136  }
   137  
   138  func ReloadASN() {
   139  	mihomoOnce.Reset(&ASNonce)
   140  }