package acl import ( "archive/tar" "compress/gzip" "io" "net/http" "net/http/httptest" "path/filepath" "strings" "testing" "time" "github.com/oschwald/maxminddb-golang" "github.com/rs/zerolog" "github.com/yusing/go-proxy/internal/task" ) func Test_dbPath(t *testing.T) { tmpDataDir := "/tmp/testdata" oldDataDir := dataDir dataDir = tmpDataDir defer func() { dataDir = oldDataDir }() tests := []struct { name string dbType MaxMindDatabaseType want string }{ {"GeoLite", MaxMindGeoLite, filepath.Join(tmpDataDir, "GeoLite2-City.mmdb")}, {"GeoIP2", MaxMindGeoIP2, filepath.Join(tmpDataDir, "GeoIP2-City.mmdb")}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { if got := dbPath(tt.dbType); got != tt.want { t.Errorf("dbPath() = %v, want %v", got, tt.want) } }) } } func Test_dbURL(t *testing.T) { tests := []struct { name string dbType MaxMindDatabaseType want string }{ {"GeoLite", MaxMindGeoLite, "https://download.maxmind.com/geoip/databases/GeoLite2-City/download?suffix=tar.gz"}, {"GeoIP2", MaxMindGeoIP2, "https://download.maxmind.com/geoip/databases/GeoIP2-City/download?suffix=tar.gz"}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { if got := dbURL(tt.dbType); got != tt.want { t.Errorf("dbURL() = %v, want %v", got, tt.want) } }) } } // --- Helper for MaxMindConfig --- type testLogger struct{ zerolog.Logger } func (testLogger) Info() *zerolog.Event { return &zerolog.Event{} } func (testLogger) Warn() *zerolog.Event { return &zerolog.Event{} } func (testLogger) Err(_ error) *zerolog.Event { return &zerolog.Event{} } func Test_MaxMindConfig_newReq(t *testing.T) { cfg := &MaxMindConfig{ AccountID: "testid", LicenseKey: "testkey", Database: MaxMindGeoLite, logger: zerolog.Nop(), } // Patch httpClient to use httptest server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if u, p, ok := r.BasicAuth(); !ok || u != "testid" || p != "testkey" { t.Errorf("basic auth not set correctly") } w.WriteHeader(http.StatusOK) })) defer server.Close() oldURL := dbURL dbURL = func(MaxMindDatabaseType) string { return server.URL } defer func() { dbURL = oldURL }() resp, err := cfg.newReq(http.MethodGet) if err != nil { t.Fatalf("newReq() error = %v", err) } if resp.StatusCode != http.StatusOK { t.Errorf("unexpected status: %v", resp.StatusCode) } } func Test_MaxMindConfig_checkUpdate(t *testing.T) { cfg := &MaxMindConfig{ AccountID: "id", LicenseKey: "key", Database: MaxMindGeoLite, logger: zerolog.Nop(), } lastMod := time.Now().UTC().Format(http.TimeFormat) buildTime := time.Now().Add(-time.Hour) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Last-Modified", lastMod) w.WriteHeader(http.StatusOK) })) defer server.Close() oldURL := dbURL dbURL = func(MaxMindDatabaseType) string { return server.URL } defer func() { dbURL = oldURL }() latest, err := cfg.checkLastest() if err != nil { t.Fatalf("checkUpdate() error = %v", err) } if latest.Equal(buildTime) { t.Errorf("expected update needed") } } type fakeReadCloser struct { firstRead bool closed bool } func (c *fakeReadCloser) Read(p []byte) (int, error) { if !c.firstRead { c.firstRead = true return strings.NewReader("FAKEMMDB").Read(p) } return 0, io.EOF } func (c *fakeReadCloser) Close() error { c.closed = true return nil } func Test_MaxMindConfig_download(t *testing.T) { cfg := &MaxMindConfig{ AccountID: "id", LicenseKey: "key", Database: MaxMindGeoLite, logger: zerolog.Nop(), } server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gz := gzip.NewWriter(w) t := tar.NewWriter(gz) t.WriteHeader(&tar.Header{ Name: dbFilename(MaxMindGeoLite), }) t.Write([]byte("1234")) t.Close() gz.Close() })) defer server.Close() oldURL := dbURL dbURL = func(MaxMindDatabaseType) string { return server.URL } defer func() { dbURL = oldURL }() tmpDir := t.TempDir() oldDataDir := dataDir dataDir = tmpDir defer func() { dataDir = oldDataDir }() // Patch maxminddb.Open to always succeed origOpen := maxmindDBOpen maxmindDBOpen = func(path string) (*maxminddb.Reader, error) { return &maxminddb.Reader{}, nil } defer func() { maxmindDBOpen = origOpen }() req, err := http.NewRequest(http.MethodGet, server.URL, nil) if err != nil { t.Fatalf("newReq() error = %v", err) } rw := httptest.NewRecorder() oldNewReq := newReq newReq = func(cfg *MaxMindConfig, method string) (*http.Response, error) { server.Config.Handler.ServeHTTP(rw, req) return rw.Result(), nil } defer func() { newReq = oldNewReq }() err = cfg.download() if err != nil { t.Fatalf("download() error = %v", err) } if cfg.db.Reader == nil { t.Error("expected db instance") } } func Test_MaxMindConfig_loadMaxMindDB(t *testing.T) { // This test should cover both the path where DB exists and where it does not // For brevity, only the non-existing path is tested here cfg := &MaxMindConfig{ AccountID: "id", LicenseKey: "key", Database: MaxMindGeoLite, logger: zerolog.Nop(), } oldOpen := maxmindDBOpen maxmindDBOpen = func(path string) (*maxminddb.Reader, error) { return &maxminddb.Reader{}, nil } defer func() { maxmindDBOpen = oldOpen }() oldDBPath := dbPath dbPath = func(MaxMindDatabaseType) string { return filepath.Join(t.TempDir(), "maxmind.mmdb") } defer func() { dbPath = oldDBPath }() task := task.RootTask("test") defer task.Finish(nil) err := cfg.LoadMaxMindDB(task) if err != nil { t.Errorf("loadMaxMindDB() error = %v", err) } }