github.com/p4gefau1t/trojan-go@v0.10.6/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/p4gefau1t/trojan-go/common" 13 "github.com/p4gefau1t/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 }