Skip to content

Commit da9092c

Browse files
PascalMinderPascalMinder
committed
Allow customizing denied request status code
1 parent 0728372 commit da9092c

3 files changed

Lines changed: 106 additions & 23 deletions

File tree

‎geoblock.go‎

Lines changed: 26 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,13 @@ import (
1616
)
1717

1818
const (
19-
xForwardedFor = "X-Forwarded-For"
20-
xRealIP = "X-Real-IP"
21-
countryHeader = "X-IPCountry"
22-
numberOfHoursInMonth = 30 * 24
23-
unknownCountryCode = "AA"
24-
countryCodeLength = 2
19+
xForwardedFor = "X-Forwarded-For"
20+
xRealIP = "X-Real-IP"
21+
countryHeader = "X-IPCountry"
22+
numberOfHoursInMonth = 30 * 24
23+
unknownCountryCode = "AA"
24+
countryCodeLength = 2
25+
defaultDeniedRequestHTTPStatusCode = 403
2526
)
2627

2728
var (
@@ -47,6 +48,7 @@ type Config struct {
4748
Countries []string `yaml:"countries,omitempty"`
4849
AllowedIPAddresses []string `yaml:"allowedIPAddresses,omitempty"`
4950
AddCountryHeader bool `yaml:"addCountryHeader"`
51+
HTTPStatusCodeDeniedRequest int `yaml:"httpStatusCodeDeniedRequest"`
5052
}
5153

5254
type ipEntry struct {
@@ -80,6 +82,7 @@ type GeoBlock struct {
8082
allowedIPRanges []*net.IPNet
8183
privateIPRanges []*net.IPNet
8284
addCountryHeader bool
85+
httpStatusCodeDeniedRequest int
8386
database *lru.LRUCache
8487
name string
8588
}
@@ -98,6 +101,15 @@ func New(_ context.Context, next http.Handler, config *Config, name string) (htt
98101
config.APITimeoutMs = 750
99102
}
100103

104+
if config.HTTPStatusCodeDeniedRequest != 0 {
105+
// check if given status code is valid
106+
if len(http.StatusText(config.HTTPStatusCodeDeniedRequest)) == 0 {
107+
return nil, fmt.Errorf("invalid denied request status code supplied")
108+
}
109+
} else {
110+
config.HTTPStatusCodeDeniedRequest = defaultDeniedRequestHTTPStatusCode
111+
}
112+
101113
var allowedIPAddresses []net.IP
102114
var allowedIPRanges []*net.IPNet
103115
for _, ipAddressEntry := range config.AllowedIPAddresses {
@@ -137,6 +149,7 @@ func New(_ context.Context, next http.Handler, config *Config, name string) (htt
137149
infoLogger.Printf("blacklist mode: %t", config.BlackListMode)
138150
infoLogger.Printf("add country header: %t", config.AddCountryHeader)
139151
infoLogger.Printf("countries: %v", config.Countries)
152+
infoLogger.Printf("Denied request status code: %d", config.HTTPStatusCodeDeniedRequest)
140153
}
141154

142155
cache, err := lru.NewLRUCache(config.CacheSize)
@@ -165,6 +178,7 @@ func New(_ context.Context, next http.Handler, config *Config, name string) (htt
165178
privateIPRanges: initPrivateIPBlocks(),
166179
database: cache,
167180
addCountryHeader: config.AddCountryHeader,
181+
httpStatusCodeDeniedRequest: config.HTTPStatusCodeDeniedRequest,
168182
name: name,
169183
}, nil
170184
}
@@ -193,7 +207,7 @@ func (a *GeoBlock) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
193207
if a.logLocalRequests {
194208
infoLogger.Println("Local ip denied: ", ipAddress)
195209
}
196-
rw.WriteHeader(http.StatusForbidden)
210+
rw.WriteHeader(a.httpStatusCodeDeniedRequest)
197211
}
198212

199213
return
@@ -223,7 +237,7 @@ func (a *GeoBlock) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
223237
entry, err = a.createNewIPEntry(req, ipAddressString)
224238

225239
if err != nil && !(os.IsTimeout(err) && a.ignoreAPITimeout) {
226-
rw.WriteHeader(http.StatusForbidden)
240+
rw.WriteHeader(a.httpStatusCodeDeniedRequest)
227241
return
228242
} else if os.IsTimeout(err) && a.ignoreAPITimeout {
229243
infoLogger.Printf("%s: request allowed [%s] due to API timeout!", a.name, ipAddress)
@@ -242,7 +256,7 @@ func (a *GeoBlock) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
242256
entry, err = a.createNewIPEntry(req, ipAddressString)
243257

244258
if err != nil {
245-
rw.WriteHeader(http.StatusForbidden)
259+
rw.WriteHeader(a.httpStatusCodeDeniedRequest)
246260
return
247261
}
248262
}
@@ -253,7 +267,7 @@ func (a *GeoBlock) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
253267

254268
if !isAllowed {
255269
infoLogger.Printf("%s: request denied [%s] for country [%s]", a.name, ipAddress, entry.Country)
256-
rw.WriteHeader(http.StatusForbidden)
270+
rw.WriteHeader(a.httpStatusCodeDeniedRequest)
257271

258272
return
259273
} else if a.logAllowedRequests {
@@ -322,7 +336,7 @@ func (a *GeoBlock) createNewIPEntry(req *http.Request, ipAddressString string) (
322336

323337
func (a *GeoBlock) getCountryCode(req *http.Request, ipAddressString string) (string, error) {
324338
if len(a.iPGeolocationHTTPHeaderField) != 0 {
325-
country, err := a.readIpGeolocationHttpHeader(req, a.iPGeolocationHTTPHeaderField)
339+
country, err := a.readIPGeolocationHTTPHeader(req, a.iPGeolocationHTTPHeaderField)
326340
if err == nil {
327341
return country, nil
328342
}
@@ -391,7 +405,7 @@ func (a *GeoBlock) callGeoJS(ipAddress string) (string, error) {
391405
return countryCode, nil
392406
}
393407

394-
func (a *GeoBlock) readIpGeolocationHttpHeader(req *http.Request, name string) (string, error) {
408+
func (a *GeoBlock) readIPGeolocationHTTPHeader(req *http.Request, name string) (string, error) {
395409
countryCode := req.Header.Get(name)
396410

397411
if len([]rune(countryCode)) != countryCodeLength {

‎geoblock_test.go‎

Lines changed: 76 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ const (
2020
invalidIP = "192.168.1.X"
2121
unknownCountry = "1.1.1.1"
2222
apiURI = "https://get.geojs.io/v1/ip/country/{ip}"
23-
ipGeolocationHttpHeaderField = "cf-ipcountry"
23+
ipGeolocationHTTPHeaderField = "cf-ipcountry"
2424
)
2525

2626
func TestEmptyApi(t *testing.T) {
@@ -35,7 +35,7 @@ func TestEmptyApi(t *testing.T) {
3535

3636
// expect error
3737
if err == nil {
38-
t.Fatal("Empty API uri accepted")
38+
t.Fatal("empty API uri accepted")
3939
}
4040
}
4141

@@ -51,7 +51,7 @@ func TestMissingIpInApi(t *testing.T) {
5151

5252
// expect error
5353
if err == nil {
54-
t.Fatal("Missing IP block in API uri")
54+
t.Fatal("missing IP block in API uri")
5555
}
5656
}
5757

@@ -65,7 +65,37 @@ func TestEmptyAllowedCountryList(t *testing.T) {
6565

6666
// expect error
6767
if err == nil {
68-
t.Fatal("Empty country list is not allowed")
68+
t.Fatal("empty country list is not allowed")
69+
}
70+
}
71+
72+
func TestEmptyDeniedRequestStatusCode(t *testing.T) {
73+
cfg := createTesterConfig()
74+
cfg.Countries = append(cfg.Countries, "CH")
75+
76+
ctx := context.Background()
77+
next := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {})
78+
79+
_, err := geoblock.New(ctx, next, cfg, "GeoBlock")
80+
81+
if err != nil {
82+
t.Fatal("no error expected for empty denied request status code")
83+
}
84+
}
85+
86+
func TestInvalidDeniedRequestStatusCode(t *testing.T) {
87+
cfg := createTesterConfig()
88+
cfg.Countries = append(cfg.Countries, "CH")
89+
cfg.HTTPStatusCodeDeniedRequest = 1
90+
91+
ctx := context.Background()
92+
next := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {})
93+
94+
_, err := geoblock.New(ctx, next, cfg, "GeoBlock")
95+
96+
// expect error
97+
if err == nil {
98+
t.Fatal("invalid denied request status code supplied")
6999
}
70100
}
71101

@@ -229,6 +259,33 @@ func TestDeniedCountry(t *testing.T) {
229259
assertStatusCode(t, recorder.Result(), http.StatusForbidden)
230260
}
231261

262+
func TestCustomDeniedRequestStatusCode(t *testing.T) {
263+
cfg := createTesterConfig()
264+
cfg.Countries = append(cfg.Countries, "CH")
265+
cfg.HTTPStatusCodeDeniedRequest = 418
266+
267+
ctx := context.Background()
268+
next := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {})
269+
270+
handler, err := geoblock.New(ctx, next, cfg, "GeoBlock")
271+
if err != nil {
272+
t.Fatal(err)
273+
}
274+
275+
recorder := httptest.NewRecorder()
276+
277+
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://localhost", nil)
278+
if err != nil {
279+
t.Fatal(err)
280+
}
281+
282+
req.Header.Add(xForwardedFor, caExampleIP)
283+
284+
handler.ServeHTTP(recorder, req)
285+
286+
assertStatusCode(t, recorder.Result(), http.StatusTeapot)
287+
}
288+
232289
func TestAllowBlacklistMode(t *testing.T) {
233290
cfg := createTesterConfig()
234291
cfg.BlackListMode = true
@@ -663,7 +720,7 @@ func TestIpGeolocationHttpField(t *testing.T) {
663720
cfg := createTesterConfig()
664721
cfg.Countries = append(cfg.Countries, "CA")
665722
cfg.AddCountryHeader = true
666-
cfg.IPGeolocationHTTPHeaderField = ipGeolocationHttpHeaderField
723+
cfg.IPGeolocationHTTPHeaderField = ipGeolocationHTTPHeaderField
667724

668725
ctx := context.Background()
669726
next := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {})
@@ -681,7 +738,7 @@ func TestIpGeolocationHttpField(t *testing.T) {
681738
}
682739

683740
req.Header.Add(xForwardedFor, caExampleIP)
684-
req.Header.Add(ipGeolocationHttpHeaderField, "CA")
741+
req.Header.Add(ipGeolocationHTTPHeaderField, "CA")
685742

686743
handler.ServeHTTP(recorder, req)
687744

@@ -698,7 +755,7 @@ func TestIpGeolocationHttpFieldContentInvalid(t *testing.T) {
698755
cfg := createTesterConfig()
699756
cfg.API = apiStub.URL + "/{ip}"
700757
cfg.Countries = append(cfg.Countries, "CA")
701-
cfg.IPGeolocationHTTPHeaderField = ipGeolocationHttpHeaderField
758+
cfg.IPGeolocationHTTPHeaderField = ipGeolocationHTTPHeaderField
702759

703760
ctx := context.Background()
704761
next := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {})
@@ -716,7 +773,7 @@ func TestIpGeolocationHttpFieldContentInvalid(t *testing.T) {
716773
}
717774

718775
req.Header.Add(xForwardedFor, caExampleIP)
719-
req.Header.Add(ipGeolocationHttpHeaderField, "")
776+
req.Header.Add(ipGeolocationHTTPHeaderField, "")
720777

721778
handler.ServeHTTP(recorder, req)
722779

@@ -745,19 +802,27 @@ type CountryCodeHandler struct {
745802

746803
func (h *CountryCodeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
747804
w.WriteHeader(http.StatusOK)
748-
w.Write([]byte(h.ResponseCountryCode))
805+
806+
_, err := w.Write([]byte(h.ResponseCountryCode))
807+
if err != nil {
808+
fmt.Println("Error on write")
809+
}
749810
}
750811

751812
func apiHandlerInvalid(w http.ResponseWriter, r *http.Request) {
752813
fmt.Fprintf(w, "Invalid Response")
753814
}
754815

755816
func apiTimeout(w http.ResponseWriter, r *http.Request) {
756-
var result = ``
757817
// Add waiting time for response
758818
time.Sleep(20 * time.Millisecond)
819+
759820
w.WriteHeader(http.StatusOK)
760-
w.Write([]byte(result))
821+
822+
_, err := w.Write([]byte(""))
823+
if err != nil {
824+
fmt.Println("Error on write")
825+
}
761826
}
762827

763828
func createTesterConfig() *geoblock.Config {

‎readme.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -510,3 +510,7 @@ allowedIPAddresses:
510510
### Add Header to request with Country Code: `addCountryHeader`
511511

512512
If set to `true`, adds the X-IPCountry header to the HTTP request header. The header contains the two letter country code returned by cache or API request.
513+
514+
### Customize denied request status code `httpStatusCodeDeniedRequest`
515+
516+
Allows customizing the HTTP status code returned if the request was denied.

0 commit comments

Comments
 (0)