github.com/ipfans/trojan-go@v0.11.0/common/geodata/decode_test.go (about)

     1  package geodata_test
     2  
     3  import (
     4  	"bytes"
     5  	"errors"
     6  	"io/fs"
     7  	"os"
     8  	"path/filepath"
     9  	"runtime"
    10  	"testing"
    11  
    12  	"github.com/ipfans/trojan-go/common"
    13  	"github.com/ipfans/trojan-go/common/geodata"
    14  )
    15  
    16  func init() {
    17  	const (
    18  		geoipURL   = "https://raw.githubusercontent.com/v2fly/geoip/release/geoip.dat"
    19  		geositeURL = "https://raw.githubusercontent.com/v2fly/domain-list-community/release/dlc.dat"
    20  	)
    21  
    22  	wd, err := os.Getwd()
    23  	common.Must(err)
    24  
    25  	tempPath := filepath.Join(wd, "..", "..", "test", "temp")
    26  	os.Setenv("TROJAN_GO_LOCATION_ASSET", tempPath)
    27  
    28  	geoipPath := common.GetAssetLocation("geoip.dat")
    29  	geositePath := common.GetAssetLocation("geosite.dat")
    30  
    31  	if _, err := os.Stat(geoipPath); err != nil && errors.Is(err, fs.ErrNotExist) {
    32  		common.Must(os.MkdirAll(tempPath, 0o755))
    33  		geoipBytes, err := common.FetchHTTPContent(geoipURL)
    34  		common.Must(err)
    35  		common.Must(common.WriteFile(geoipPath, geoipBytes))
    36  	}
    37  	if _, err := os.Stat(geositePath); err != nil && errors.Is(err, fs.ErrNotExist) {
    38  		common.Must(os.MkdirAll(tempPath, 0o755))
    39  		geositeBytes, err := common.FetchHTTPContent(geositeURL)
    40  		common.Must(err)
    41  		common.Must(common.WriteFile(geositePath, geositeBytes))
    42  	}
    43  }
    44  
    45  func TestDecodeGeoIP(t *testing.T) {
    46  	filename := common.GetAssetLocation("geoip.dat")
    47  	result, err := geodata.Decode(filename, "test")
    48  	if err != nil {
    49  		t.Error(err)
    50  	}
    51  
    52  	expected := []byte{10, 4, 84, 69, 83, 84, 18, 8, 10, 4, 127, 0, 0, 0, 16, 8}
    53  	if !bytes.Equal(result, expected) {
    54  		t.Errorf("failed to load geoip:test, expected: %v, got: %v", expected, result)
    55  	}
    56  }
    57  
    58  func TestDecodeGeoSite(t *testing.T) {
    59  	filename := common.GetAssetLocation("geosite.dat")
    60  	result, err := geodata.Decode(filename, "test")
    61  	if err != nil {
    62  		t.Error(err)
    63  	}
    64  
    65  	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}
    66  	if !bytes.Equal(result, expected) {
    67  		t.Errorf("failed to load geosite:test, expected: %v, got: %v", expected, result)
    68  	}
    69  }
    70  
    71  func BenchmarkLoadGeoIP(b *testing.B) {
    72  	m1 := runtime.MemStats{}
    73  	m2 := runtime.MemStats{}
    74  
    75  	loader := geodata.NewGeodataLoader()
    76  
    77  	runtime.ReadMemStats(&m1)
    78  	cn, _ := loader.LoadGeoIP("cn")
    79  	private, _ := loader.LoadGeoIP("private")
    80  	runtime.KeepAlive(cn)
    81  	runtime.KeepAlive(private)
    82  	runtime.ReadMemStats(&m2)
    83  
    84  	b.ReportMetric(float64(m2.Alloc-m1.Alloc)/1024, "KiB(GeoIP-Alloc)")
    85  	b.ReportMetric(float64(m2.TotalAlloc-m1.TotalAlloc)/1024/1024, "MiB(GeoIP-TotalAlloc)")
    86  }
    87  
    88  func BenchmarkLoadGeoSite(b *testing.B) {
    89  	m3 := runtime.MemStats{}
    90  	m4 := runtime.MemStats{}
    91  
    92  	loader := geodata.NewGeodataLoader()
    93  
    94  	runtime.ReadMemStats(&m3)
    95  	cn, _ := loader.LoadGeoSite("cn")
    96  	notcn, _ := loader.LoadGeoSite("geolocation-!cn")
    97  	private, _ := loader.LoadGeoSite("private")
    98  	runtime.KeepAlive(cn)
    99  	runtime.KeepAlive(notcn)
   100  	runtime.KeepAlive(private)
   101  	runtime.ReadMemStats(&m4)
   102  
   103  	b.ReportMetric(float64(m4.Alloc-m3.Alloc)/1024/1024, "MiB(GeoSite-Alloc)")
   104  	b.ReportMetric(float64(m4.TotalAlloc-m3.TotalAlloc)/1024/1024, "MiB(GeoSite-TotalAlloc)")
   105  }