package service import ( "errors" "fmt" "net/http" "net/http/httptest" "os" "path/filepath" "testing" "github.com/mhsanaei/3x-ui/v2/database/model" ) func TestNormalizeAliasKey(t *testing.T) { if got := NormalizeAliasKey("GeoIP-IR"); got != "geoip_ir" { t.Fatalf("got %q", got) } if got := NormalizeAliasKey("a-b_c"); got != "a_b_c" { t.Fatalf("got %q", got) } } func TestNewCustomGeoService(t *testing.T) { s := NewCustomGeoService() if err := s.validateAlias("ok_alias-1"); err != nil { t.Fatal(err) } } func TestTriggerUpdateAllAllSuccess(t *testing.T) { s := CustomGeoService{} s.updateAllGetAll = func() ([]model.CustomGeoResource, error) { return []model.CustomGeoResource{ {Id: 1, Alias: "a"}, {Id: 2, Alias: "b"}, }, nil } s.updateAllApply = func(id int, onStartup bool) (string, error) { return fmt.Sprintf("geo_%d.dat", id), nil } restartCalls := 0 s.updateAllRestart = func() error { restartCalls++ return nil } res, err := s.TriggerUpdateAll() if err != nil { t.Fatal(err) } if len(res.Succeeded) != 2 || len(res.Failed) != 0 { t.Fatalf("unexpected result: %+v", res) } if restartCalls != 1 { t.Fatalf("expected 1 restart, got %d", restartCalls) } } func TestTriggerUpdateAllPartialSuccess(t *testing.T) { s := CustomGeoService{} s.updateAllGetAll = func() ([]model.CustomGeoResource, error) { return []model.CustomGeoResource{ {Id: 1, Alias: "ok"}, {Id: 2, Alias: "bad"}, }, nil } s.updateAllApply = func(id int, onStartup bool) (string, error) { if id == 2 { return "geo_2.dat", ErrCustomGeoDownload } return "geo_1.dat", nil } restartCalls := 0 s.updateAllRestart = func() error { restartCalls++ return nil } res, err := s.TriggerUpdateAll() if err != nil { t.Fatal(err) } if len(res.Succeeded) != 1 || len(res.Failed) != 1 { t.Fatalf("unexpected result: %+v", res) } if restartCalls != 1 { t.Fatalf("expected 1 restart, got %d", restartCalls) } } func TestTriggerUpdateAllAllFailure(t *testing.T) { s := CustomGeoService{} s.updateAllGetAll = func() ([]model.CustomGeoResource, error) { return []model.CustomGeoResource{ {Id: 1, Alias: "a"}, {Id: 2, Alias: "b"}, }, nil } s.updateAllApply = func(id int, onStartup bool) (string, error) { return fmt.Sprintf("geo_%d.dat", id), ErrCustomGeoDownload } restartCalls := 0 s.updateAllRestart = func() error { restartCalls++ return nil } res, err := s.TriggerUpdateAll() if err != nil { t.Fatal(err) } if len(res.Succeeded) != 0 || len(res.Failed) != 2 { t.Fatalf("unexpected result: %+v", res) } if restartCalls != 0 { t.Fatalf("expected 0 restart, got %d", restartCalls) } } func TestCustomGeoValidateAlias(t *testing.T) { s := CustomGeoService{} if err := s.validateAlias(""); !errors.Is(err, ErrCustomGeoAliasRequired) { t.Fatal("empty alias") } if err := s.validateAlias("Bad"); !errors.Is(err, ErrCustomGeoAliasPattern) { t.Fatal("uppercase") } if err := s.validateAlias("a b"); !errors.Is(err, ErrCustomGeoAliasPattern) { t.Fatal("space") } if err := s.validateAlias("ok_alias-1"); err != nil { t.Fatal(err) } if err := s.validateAlias("geoip"); !errors.Is(err, ErrCustomGeoAliasReserved) { t.Fatal("reserved") } } func TestCustomGeoValidateURL(t *testing.T) { s := CustomGeoService{} if err := s.validateURL(""); !errors.Is(err, ErrCustomGeoURLRequired) { t.Fatal("empty") } if err := s.validateURL("ftp://x"); !errors.Is(err, ErrCustomGeoURLScheme) { t.Fatal("ftp") } if err := s.validateURL("https://example.com/a.dat"); err != nil { t.Fatal(err) } } func TestCustomGeoValidateType(t *testing.T) { s := CustomGeoService{} if err := s.validateType("geosite"); err != nil { t.Fatal(err) } if err := s.validateType("x"); !errors.Is(err, ErrCustomGeoInvalidType) { t.Fatal("bad type") } } func TestCustomGeoDownloadToPath(t *testing.T) { ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("X-Test", "1") if r.Header.Get("If-Modified-Since") != "" { w.WriteHeader(http.StatusNotModified) return } w.WriteHeader(http.StatusOK) _, _ = w.Write(make([]byte, minDatBytes+1)) })) defer ts.Close() dir := t.TempDir() t.Setenv("XUI_BIN_FOLDER", dir) dest := filepath.Join(dir, "geoip_t.dat") s := CustomGeoService{} skipped, _, err := s.downloadToPath(ts.URL, dest, "") if err != nil { t.Fatal(err) } if skipped { t.Fatal("expected download") } st, err := os.Stat(dest) if err != nil || st.Size() < minDatBytes { t.Fatalf("file %v", err) } skipped2, _, err2 := s.downloadToPath(ts.URL, dest, "") if err2 != nil || !skipped2 { t.Fatalf("304 expected skipped=%v err=%v", skipped2, err2) } } func TestCustomGeoDownloadToPath_missingLocalSendsNoIMSFromDB(t *testing.T) { lm := "Wed, 21 Oct 2015 07:28:00 GMT" ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Header.Get("If-Modified-Since") != "" { w.WriteHeader(http.StatusNotModified) return } w.Header().Set("Last-Modified", lm) w.WriteHeader(http.StatusOK) _, _ = w.Write(make([]byte, minDatBytes+1)) })) defer ts.Close() dir := t.TempDir() t.Setenv("XUI_BIN_FOLDER", dir) dest := filepath.Join(dir, "geoip_rebuild.dat") s := CustomGeoService{} skipped, _, err := s.downloadToPath(ts.URL, dest, lm) if err != nil { t.Fatal(err) } if skipped { t.Fatal("must not treat as not-modified when local file is missing") } if _, err := os.Stat(dest); err != nil { t.Fatal("file should exist after container-style rebuild") } } func TestCustomGeoDownloadToPath_repairSkipsConditional(t *testing.T) { lm := "Wed, 21 Oct 2015 07:28:00 GMT" ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Header.Get("If-Modified-Since") != "" { w.WriteHeader(http.StatusNotModified) return } w.Header().Set("Last-Modified", lm) w.WriteHeader(http.StatusOK) _, _ = w.Write(make([]byte, minDatBytes+1)) })) defer ts.Close() dir := t.TempDir() t.Setenv("XUI_BIN_FOLDER", dir) dest := filepath.Join(dir, "geoip_bad.dat") if err := os.WriteFile(dest, make([]byte, minDatBytes-1), 0o644); err != nil { t.Fatal(err) } s := CustomGeoService{} skipped, _, err := s.downloadToPath(ts.URL, dest, lm) if err != nil { t.Fatal(err) } if skipped { t.Fatal("corrupt local file must be re-downloaded, not 304") } st, err := os.Stat(dest) if err != nil || st.Size() < minDatBytes { t.Fatalf("file repaired: %v", err) } } func TestCustomGeoFileNameFor(t *testing.T) { s := CustomGeoService{} if s.fileNameFor("geoip", "a") != "geoip_a.dat" { t.Fatal("geoip name") } if s.fileNameFor("geosite", "b") != "geosite_b.dat" { t.Fatal("geosite name") } } func TestLocalDatFileNeedsRepair(t *testing.T) { dir := t.TempDir() if !localDatFileNeedsRepair(filepath.Join(dir, "missing.dat")) { t.Fatal("missing") } smallPath := filepath.Join(dir, "small.dat") if err := os.WriteFile(smallPath, make([]byte, minDatBytes-1), 0o644); err != nil { t.Fatal(err) } if !localDatFileNeedsRepair(smallPath) { t.Fatal("small") } okPath := filepath.Join(dir, "ok.dat") if err := os.WriteFile(okPath, make([]byte, minDatBytes), 0o644); err != nil { t.Fatal(err) } if localDatFileNeedsRepair(okPath) { t.Fatal("ok size") } dirPath := filepath.Join(dir, "isdir.dat") if err := os.Mkdir(dirPath, 0o755); err != nil { t.Fatal(err) } if !localDatFileNeedsRepair(dirPath) { t.Fatal("dir should need repair") } if !CustomGeoLocalFileNeedsRepair(dirPath) { t.Fatal("exported wrapper dir") } if CustomGeoLocalFileNeedsRepair(okPath) { t.Fatal("exported wrapper ok file") } } func TestProbeCustomGeoURL_HEADOK(t *testing.T) { ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method == http.MethodHead { w.WriteHeader(http.StatusOK) return } w.WriteHeader(http.StatusOK) })) defer ts.Close() if err := probeCustomGeoURL(ts.URL); err != nil { t.Fatal(err) } } func TestProbeCustomGeoURL_HEAD405GETRange(t *testing.T) { ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method == http.MethodHead { w.WriteHeader(http.StatusMethodNotAllowed) return } if r.Method == http.MethodGet && r.Header.Get("Range") != "" { w.WriteHeader(http.StatusPartialContent) _, _ = w.Write([]byte{0}) return } w.WriteHeader(http.StatusBadRequest) })) defer ts.Close() if err := probeCustomGeoURL(ts.URL); err != nil { t.Fatal(err) } }