| 1 | package main |
| 2 | |
| 3 | import ( |
| 4 | "bytes" |
| 5 | "context" |
| 6 | "encoding/json" |
| 7 | "io" |
| 8 | "net/http" |
| 9 | "net/http/httptest" |
| 10 | "os" |
| 11 | "path/filepath" |
| 12 | "strings" |
| 13 | "sync/atomic" |
| 14 | "testing" |
| 15 | "time" |
| 16 | ) |
| 17 | |
| 18 | func TestSystemVerbKnowsGeoip(t *testing.T) { |
| 19 | if verb, rest, _ := systemVerb([]string{"geoip", "publish", "x.mmdb"}); verb != "geoip" || len(rest) != 2 { |
| 20 | t.Fatalf("systemVerb = %q %v", verb, rest) |
| 21 | } |
| 22 | } |
| 23 | |
| 24 | func TestPublishGeoipStreamsARawBodyAndFollowsTheJob(t *testing.T) { |
| 25 | geoipPollInterval = time.Millisecond |
| 26 | t.Cleanup(func() { geoipPollInterval = 2 * time.Second }) |
| 27 | db := bytes.Repeat([]byte("m"), 4096) |
| 28 | path := filepath.Join(t.TempDir(), "db.mmdb") |
| 29 | if err := os.WriteFile(path, db, 0o600); err != nil { |
| 30 | t.Fatal(err) |
| 31 | } |
| 32 | var polls atomic.Int32 |
| 33 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 34 | switch { |
| 35 | case r.Method == http.MethodPut && r.URL.Path == "/api/cli/system/geoip": |
| 36 | if r.Header.Get("Content-Type") != "application/octet-stream" || r.ContentLength != int64(len(db)) { |
| 37 | http.Error(w, "bad upload headers", http.StatusBadRequest) |
| 38 | return |
| 39 | } |
| 40 | if r.URL.Query().Get("force") != "1" { |
| 41 | http.Error(w, "force not passed", http.StatusBadRequest) |
| 42 | return |
| 43 | } |
| 44 | b, _ := io.ReadAll(r.Body) |
| 45 | if !bytes.Equal(b, db) { |
| 46 | http.Error(w, "body mangled", http.StatusBadRequest) |
| 47 | return |
| 48 | } |
| 49 | w.WriteHeader(http.StatusAccepted) |
| 50 | w.Write([]byte(`{"job":"j1"}`)) |
| 51 | case r.URL.Path == "/api/cli/system/jobs/j1": |
| 52 | if polls.Add(1) < 2 { |
| 53 | w.Write([]byte(`{"steps":["received"],"done":false,"error":""}`)) |
| 54 | return |
| 55 | } |
| 56 | w.Write([]byte(`{"steps":["received","verifying"],"done":true,"error":"","result":{"sha256":"9f2c0000aaaa","database_type":"DBIP-City-Lite","build_epoch":1790812800}}`)) |
| 57 | default: |
| 58 | http.NotFound(w, r) |
| 59 | } |
| 60 | })) |
| 61 | t.Cleanup(ts.Close) |
| 62 | sess := &apiSession{c: newConsoleClient(ts.URL, "tok"), console: ts.URL} |
| 63 | var out bytes.Buffer |
| 64 | ptr, err := sess.publishGeoip(context.Background(), path, true, &out) |
| 65 | if err != nil { |
| 66 | t.Fatal(err) |
| 67 | } |
| 68 | if ptr.DatabaseType != "DBIP-City-Lite" || ptr.SHA256 != "9f2c0000aaaa" { |
| 69 | t.Fatalf("pointer = %+v", ptr) |
| 70 | } |
| 71 | if !strings.Contains(out.String(), "verifying") { |
| 72 | t.Fatalf("steps not printed: %q", out.String()) |
| 73 | } |
| 74 | } |
| 75 | |
| 76 | func TestPublishGeoipReportsTheJobsRefusal(t *testing.T) { |
| 77 | geoipPollInterval = time.Millisecond |
| 78 | t.Cleanup(func() { geoipPollInterval = 2 * time.Second }) |
| 79 | path := filepath.Join(t.TempDir(), "db.mmdb") |
| 80 | os.WriteFile(path, []byte("junk"), 0o600) |
| 81 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 82 | if r.Method == http.MethodPut { |
| 83 | io.Copy(io.Discard, r.Body) |
| 84 | w.WriteHeader(http.StatusAccepted) |
| 85 | w.Write([]byte(`{"job":"j2"}`)) |
| 86 | return |
| 87 | } |
| 88 | json.NewEncoder(w).Encode(map[string]any{"steps": []string{}, "done": true, "error": "geoip database refused: not a valid MaxMind-format database"}) |
| 89 | })) |
| 90 | t.Cleanup(ts.Close) |
| 91 | sess := &apiSession{c: newConsoleClient(ts.URL, "tok"), console: ts.URL} |
| 92 | if _, err := sess.publishGeoip(context.Background(), path, false, io.Discard); err == nil || !strings.Contains(err.Error(), "not a valid MaxMind-format database") { |
| 93 | t.Fatalf("err = %v", err) |
| 94 | } |
| 95 | } |
| 96 | |
| 97 | func TestPrintGeoipStatusVerdicts(t *testing.T) { |
| 98 | var resp geoipStatusResponse |
| 99 | json.Unmarshal([]byte(`{"pointer":{"sha256":"9f2c0000aaaa","database_type":"GeoLite2-City","build":"2026-10-01T00:00:00Z","origin":"manual","expires_at":"2026-10-31T00:00:00Z","attribution":"This product includes GeoLite2 data"}, |
| 100 | "boxes":[{"label":"ie-1","state":"ok"},{"label":"node-1","fleet":"tito","state":"old binary"}], |
| 101 | "ready":false,"blocking":["tito/node-1: runs a binary that does not strip forged geo headers"]}`), &resp) |
| 102 | var out bytes.Buffer |
| 103 | printGeoipStatus(&out, resp, time.Date(2026, 10, 6, 0, 0, 0, 0, time.UTC)) |
| 104 | s := out.String() |
| 105 | for _, want := range []string{"GeoLite2-City", "5 days ago", "expires 2026-10-31", "tito/node-1", "not ready", "1 in the way"} { |
| 106 | if !strings.Contains(s, want) { |
| 107 | t.Errorf("status lacks %q:\n%s", want, s) |
| 108 | } |
| 109 | } |
| 110 | out.Reset() |
| 111 | printGeoipStatus(&out, geoipStatusResponse{Ready: true}, time.Now()) |
| 112 | if !strings.Contains(out.String(), "none published") || !strings.Contains(out.String(), "ready") { |
| 113 | t.Errorf("empty status = %q", out.String()) |
| 114 | } |
| 115 | } |
| 116 | |
| 117 | func TestPrintGeoipOffWarnsAboutTheFeed(t *testing.T) { |
| 118 | var out bytes.Buffer |
| 119 | printGeoipOff(&out, true) |
| 120 | if !strings.Contains(out.String(), "unset CONSOLE_GEOIP_FEED") { |
| 121 | t.Fatalf("off with the feed set = %q", out.String()) |
| 122 | } |
| 123 | out.Reset() |
| 124 | printGeoipOff(&out, false) |
| 125 | if strings.Contains(out.String(), "CONSOLE_GEOIP_FEED") { |
| 126 | t.Fatalf("off without a feed mentioned it: %q", out.String()) |
| 127 | } |
| 128 | } |
| 129 | |
| 130 | func TestPrintGeoipStatusFeedLine(t *testing.T) { |
| 131 | var resp geoipStatusResponse |
| 132 | json.Unmarshal([]byte(`{"feed":{"name":"dbip-city-lite","failures":3,"last_failure":"download"},"ready":true}`), &resp) |
| 133 | var out bytes.Buffer |
| 134 | printGeoipStatus(&out, resp, time.Now()) |
| 135 | if !strings.Contains(out.String(), "feed: dbip-city-lite, failing (3 in a row, last: download)") { |
| 136 | t.Fatalf("feed line = %q", out.String()) |
| 137 | } |
| 138 | json.Unmarshal([]byte(`{"feed":{"name":"dbip-city-lite","last_success":"2026-10-01T00:00:00Z"},"ready":true}`), &resp) |
| 139 | resp.Feed.Failures = 0 |
| 140 | out.Reset() |
| 141 | printGeoipStatus(&out, resp, time.Now()) |
| 142 | if !strings.Contains(out.String(), "feed: dbip-city-lite, last success 2026-10-01") { |
| 143 | t.Fatalf("feed line = %q", out.String()) |
| 144 | } |
| 145 | } |
| 146 | |
amadan