GeoCIDR/examples_test.go at main · F0gC10ud/GeoCIDR · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
package geocidr

import (
    "context"
    "net/http"
    "net/http/httptest"
    "os"
    "testing"
    "time"
)

// Helpers
func newMiddlewareForTest(t *testing.T, cfg *Config) *GeoCIDR {
    t.Helper()
    ctx, cancel := context.WithCancel(context.Background())
    t.Cleanup(cancel)
    h, err := New(ctx, okHandler{}, cfg, "test")
    if err != nil { t.Fatalf("New: %v", err) }
    return h.(*GeoCIDR)
}

func TestAllowPassExample(t *testing.T) {
    cfg := CreateConfig()
    cfg.MaltrailURL = "" // disable remote to isolate
    cfg.DefaultBlocklistPath = ""
    cfg.WhitelistPath = ""
    cfg.UpdateInterval = time.Hour

    mw := newMiddlewareForTest(t, cfg)
    rr := doReq(mw, "203.0.113.10", cfg)
    if rr.Code != http.StatusOK {
        t.Fatalf("expected pass: 200, got %d", rr.Code)
    }
}

func TestBlockedByGeoExample(t *testing.T) {
    cfg := CreateConfig()
    cfg.BlockCountries = []string{"RU"}
    cfg.MaltrailURL = ""
    cfg.UpdateInterval = time.Hour
    mw := newMiddlewareForTest(t, cfg)
    // inject fake geo
    mw.geo = fakeGeo{m: map[string]string{"203.0.113.5": "RU"}}
    rr := doReq(mw, "203.0.113.5", cfg)
    if rr.Code != http.StatusForbidden {
        t.Fatalf("expected geo block: 403, got %d", rr.Code)
    }
}

func TestBlockedByDefaultBlocklistExample(t *testing.T) {
    dir := t.TempDir()
    def := dir + "/default.txt"
    if err := os.WriteFile(def, []byte("10.0.0.0/8\n"), 0o600); err != nil {
        t.Fatal(err)
    }
    cfg := CreateConfig()
    cfg.DefaultBlocklistPath = def
    cfg.MaltrailURL = ""
    cfg.UpdateInterval = 10 * time.Millisecond
    mw := newMiddlewareForTest(t, cfg)
    // allow initial load
    time.Sleep(30 * time.Millisecond)
    rr := doReq(mw, "10.1.2.3", cfg)
    if rr.Code != http.StatusForbidden {
        t.Fatalf("expected default list block: 403, got %d", rr.Code)
    }
}

func TestDynamicReload_DefaultBlocklist(t *testing.T) {
    dir := t.TempDir()
    def := dir + "/default.txt"
    // initial default: block 10.0.0.0/8
    if err := os.WriteFile(def, []byte("10.0.0.0/8\n"), 0o600); err != nil { t.Fatal(err) }
    cfg := CreateConfig()
    cfg.DefaultBlocklistPath = def
    cfg.MaltrailURL = "" // isolate
    cfg.UpdateInterval = time.Hour // disable periodic update churn
    cfg.FileWatchInterval = 10 * time.Millisecond
    mw := newMiddlewareForTest(t, cfg)
    time.Sleep(30 * time.Millisecond) // allow initial load

    // Initially blocked
    rr := doReq(mw, "10.1.2.3", cfg)
    if rr.Code != http.StatusForbidden { t.Fatalf("expected 403 initially, got %d", rr.Code) }

    // Update default list to block a different CIDR
    if err := os.WriteFile(def, []byte("203.0.113.0/24\n"), 0o600); err != nil { t.Fatal(err) }
    // Wait for watcher to detect and reload
    time.Sleep(40 * time.Millisecond)

    // Old IP should now pass
    rr = doReq(mw, "10.1.2.3", cfg)
    if rr.Code != http.StatusOK { t.Fatalf("expected 200 after reload for 10.1.2.3, got %d", rr.Code) }
    // New CIDR should be blocked
    rr = doReq(mw, "203.0.113.10", cfg)
    if rr.Code != http.StatusForbidden { t.Fatalf("expected 403 after reload for 203.0.113.10, got %d", rr.Code) }
}

func TestBlockedByMaltrailBlocklistExample(t *testing.T) {
    // Requires HTTP client/server; skip under restricted CI sandboxes.
    skipServerTests(t)
    // Serve a small maltrail-like list
    ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        w.WriteHeader(200)
        _, _ = w.Write([]byte("5.6.7.0/24\n"))
    }))
    defer ts.Close()
    cfg := CreateConfig()
    cfg.MaltrailURL = ts.URL
    cfg.UpdateInterval = 10 * time.Millisecond
    mw := newMiddlewareForTest(t, cfg)
    // allow initial load
    time.Sleep(30 * time.Millisecond)
    rr := doReq(mw, "5.6.7.8", cfg)
    if rr.Code != http.StatusForbidden {
        t.Fatalf("expected maltrail block: 403, got %d", rr.Code)
    }
}