headerblock/headerblock_test.go at master · LineTW/headerblock · 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
119
120
121
122
123
124
125
126
127
128
129
130
131
package headerblock_test

import (
	"context"
	"log"
	"net/http"
	"net/http/httptest"
	"testing"

	tbua "github.com/nhomchatgpt/headerblock"
)

const pluginName = "headerBlock"

type noopHandler struct{}

func (n noopHandler) ServeHTTP(rw http.ResponseWriter, _ *http.Request) {
	rw.WriteHeader(http.StatusTeapot)
}

func TestPlugin(t *testing.T) {
	// Test for requests without any User-Agent headers
	t.Run("NoUserAgents", func(t *testing.T) {
		cfg := tbua.CreateConfig()
		cfg.Log = true
		p, err := tbua.New(context.Background(), noopHandler{}, cfg, pluginName)
		if err != nil {
			t.Fatalf("unexpected error during plugin creation: %v", err)
		}

		req := httptest.NewRequest(http.MethodGet, "/foobar", nil)
		rr := httptest.NewRecorder()
		p.ServeHTTP(rr, req)

		if rr.Code != http.StatusTeapot {
			t.Fatalf("unexpected status: got %v, expected %v", rr.Code, http.StatusTeapot)
		}
	})

	// Test for requests with a non-blocked User-Agent header
	t.Run("ValidUserAgent", func(t *testing.T) {
		cfg := tbua.CreateConfig()
		cfg.RequestHeaders = append(cfg.RequestHeaders, tbua.HeaderConfig{
			Name:  "User-Agent",
			Value: "SpamBot",
		})

		p, err := tbua.New(context.Background(), noopHandler{}, cfg, pluginName)
		if err != nil {
			t.Fatalf("unexpected error during plugin creation: %v", err)
		}

		req := httptest.NewRequest(http.MethodGet, "/foobar", nil)
		req.Header.Set("User-Agent", "ValidUserAgent")
		rr := httptest.NewRecorder()
		p.ServeHTTP(rr, req)

		if rr.Code != http.StatusTeapot {
			t.Fatalf("unexpected status: got %v, expected %v for header %v", rr.Code, http.StatusTeapot, req.Header.Get("User-Agent"))
		}
	})

	// Test for requests with a blocked User-Agent header
	t.Run("ForbiddenUserAgent", func(t *testing.T) {
		cfg := tbua.CreateConfig()
		cfg.RequestHeaders = append(cfg.RequestHeaders, tbua.HeaderConfig{
			Name:  "User-Agent",
			Value: "Googlebot",
		})

		p, err := tbua.New(context.Background(), noopHandler{}, cfg, pluginName)
		if err != nil {
			t.Fatalf("unexpected error during plugin creation: %v", err)
		}

		req := httptest.NewRequest(http.MethodGet, "/foobar", nil)
		req.Header.Set("User-Agent", "Googlebot")
		rr := httptest.NewRecorder()
		p.ServeHTTP(rr, req)

		if rr.Code != http.StatusForbidden {
			t.Fatalf("unexpected status: got %v, expected %v for header %v", rr.Code, http.StatusForbidden, req.Header.Get("User-Agent"))
		}
	})

	// Test for requests with CF-IPCountry whitelist
	t.Run("WhitelistCfIpCountry", func(t *testing.T) {
		cfg := tbua.CreateConfig()
		cfg.WhitelistRequestHeaders = append(cfg.WhitelistRequestHeaders, tbua.HeaderConfig{
			Name:  "Cf-Ipcountry",
			Value: "VN",
		})

		p, err := tbua.New(context.Background(), noopHandler{}, cfg, pluginName)
		if err != nil {
			t.Fatalf("unexpected error during plugin creation: %v", err)
		}

		// Allowed: CF-IPCountry is VN
		reqVN := httptest.NewRequest(http.MethodGet, "/foobar", nil)
		reqVN.Header.Set("Cf-Ipcountry", "VN")
		rrVN := httptest.NewRecorder()

		log.Printf("Request Headers: %+v", reqVN.Header) // Log request headers
		p.ServeHTTP(rrVN, reqVN)

		if rrVN.Code != http.StatusTeapot {
			t.Fatalf("unexpected status: got %v, expected %v for header %v", rrVN.Code, http.StatusTeapot, reqVN.Header.Get("CF-IPCountry"))
		}

		// Blocked: CF-IPCountry is FR
		reqFR := httptest.NewRequest(http.MethodGet, "/foobar", nil)
		reqFR.Header.Set("Cf-Ipcountry", "FR")
		rrFR := httptest.NewRecorder()
		p.ServeHTTP(rrFR, reqFR)

		if rrFR.Code != http.StatusForbidden {
			t.Fatalf("unexpected status: got %v, expected %v for header %v", rrFR.Code, http.StatusForbidden, reqFR.Header.Get("CF-IPCountry"))
		}

		// Test request with no headers
		req := httptest.NewRequest(http.MethodGet, "/foobar", nil)
		rr := httptest.NewRecorder()
		p.ServeHTTP(rr, req)

		// Expect 403 Forbidden for requests with no headers
		if rr.Code != http.StatusForbidden {
			t.Fatalf("unexpected status: got %v, expected %v for request with no headers", rr.Code, http.StatusTeapot)
		}
	})
}