add returning X-IPCountry header for explicitly allowed IP addresses (fixes #76)

Author: PascalMinder

geoblock.go

@@ -229,6 +229,12 @@ func (a *GeoBlock) allowDenyIPAddress(requestIPAddr *net.IP, req *http.Request)
 
 	// check if the request IP address is explicitly allowed
 	if ipInSlice(*requestIPAddr, a.allowedIPAddresses) {
+		if a.addCountryHeader {
+			ok, countryCode := a.cachedRequestIP(requestIPAddr, req)
+			if ok && len(countryCode) > 0 {
+				req.Header.Set(countryHeader, countryCode)
+			}
+		}
 		if a.logAllowedRequests {
 			a.infoLogger.Printf("%s: request allowed [%s] since the IP address is explicitly allowed", a.name, requestIPAddr)
 		}
@@ -238,6 +244,12 @@ func (a *GeoBlock) allowDenyIPAddress(requestIPAddr *net.IP, req *http.Request)
 	// check if the request IP address is contained within one of the explicitly allowed IP address ranges
 	for _, ipRange := range a.allowedIPRanges {
 		if ipRange.Contains(*requestIPAddr) {
+			if a.addCountryHeader {
+				ok, countryCode := a.cachedRequestIP(requestIPAddr, req)
+				if ok && len(countryCode) > 0 {
+					req.Header.Set(countryHeader, countryCode)
+				}
+			}
 			if a.logLocalRequests {
 				a.infoLogger.Printf("%s: request allowed [%s] since the IP address is explicitly allowed", a.name, requestIPAddr)
 			}
@@ -281,7 +293,6 @@ func (a *GeoBlock) allowDenyCachedRequestIP(requestIPAddr *net.IP, req *http.Req
 	// check if existing entry was made more than a month ago, if so update the entry
 	if time.Since(entry.Timestamp).Hours() >= numberOfHoursInMonth && a.forceMonthlyUpdate {
 		entry, err = a.createNewIPEntry(req, ipAddressString)
-
 		if err != nil {
 			return false, ""
 		}
@@ -303,6 +314,36 @@ func (a *GeoBlock) allowDenyCachedRequestIP(requestIPAddr *net.IP, req *http.Req
 	return true, entry.Country
 }
 
+func (a *GeoBlock) cachedRequestIP(requestIPAddr *net.IP, req *http.Request) (bool, string) {
+	ipAddressString := requestIPAddr.String()
+	cacheEntry, ok := a.database.Get(ipAddressString)
+
+	var entry ipEntry
+	var err error
+	if !ok {
+		entry, err = a.createNewIPEntry(req, ipAddressString)
+		if err != nil {
+			return false, ""
+		}
+	} else {
+		entry = cacheEntry.(ipEntry)
+	}
+
+	if a.logAPIRequests {
+		a.infoLogger.Println("Loaded from database: ", entry)
+	}
+
+	// check if existing entry was made more than a month ago, if so update the entry
+	if time.Since(entry.Timestamp).Hours() >= numberOfHoursInMonth && a.forceMonthlyUpdate {
+		entry, err = a.createNewIPEntry(req, ipAddressString)
+		if err != nil {
+			return false, ""
+		}
+	}
+
+	return true, entry.Country
+}
+
 func (a *GeoBlock) collectRemoteIP(req *http.Request) ([]*net.IP, error) {
 	var ipList []*net.IP
 

geoblock_test.go

@@ -737,7 +737,6 @@ func TestInvalidApiResponse(t *testing.T) {
 	var apiStub = httptest.NewServer(http.HandlerFunc(apiHandlerInvalid))
 
 	cfg := createTesterConfig()
-	fmt.Println(apiStub.URL)
 	cfg.API = apiStub.URL + "/{ip}"
 	cfg.Countries = append(cfg.Countries, "CH")
 
@@ -770,7 +769,6 @@ func TestApiResponseTimeoutAllowed(t *testing.T) {
 	var apiStub = httptest.NewServer(http.HandlerFunc(apiTimeout))
 
 	cfg := createTesterConfig()
-	fmt.Println(apiStub.URL)
 	cfg.API = apiStub.URL + "/{ip}"
 	cfg.Countries = append(cfg.Countries, "CH")
 	cfg.APITimeoutMs = 5
@@ -805,7 +803,6 @@ func TestApiResponseTimeoutNotAllowed(t *testing.T) {
 	var apiStub = httptest.NewServer(http.HandlerFunc(apiTimeout))
 
 	cfg := createTesterConfig()
-	fmt.Println(apiStub.URL)
 	cfg.API = apiStub.URL + "/{ip}"
 	cfg.Countries = append(cfg.Countries, "CH")
 	cfg.APITimeoutMs = 5
@@ -863,6 +860,41 @@ func TestExplicitlyAllowedIP(t *testing.T) {
 	assertStatusCode(t, recorder.Result(), http.StatusOK)
 }
 
+func TestExplicitlyAllowedIPWithIPCountryHeader(t *testing.T) {
+	// set up our fake api server
+	apiHandler := &CountryCodeHandler{ResponseCountryCode: "CA"}
+	var apiStub = httptest.NewServer(apiHandler)
+
+	cfg := createTesterConfig()
+	cfg.API = apiStub.URL + "/{ip}"
+	cfg.Countries = append(cfg.Countries, "CH")
+	cfg.AllowedIPAddresses = append(cfg.AllowedIPAddresses, caExampleIP)
+	cfg.LogLocalRequests = true
+	cfg.AddCountryHeader = true
+
+	ctx := context.Background()
+	next := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {})
+
+	handler, err := geoblock.New(ctx, next, cfg, "GeoBlock")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	recorder := httptest.NewRecorder()
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://localhost", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	req.Header.Add(xForwardedFor, caExampleIP)
+
+	handler.ServeHTTP(recorder, req)
+
+	assertStatusCode(t, recorder.Result(), http.StatusOK)
+	assertRequestHeader(t, req, CountryHeader, "CA")
+}
+
 func TestExplicitlyAllowedIPNoMatch(t *testing.T) {
 	cfg := createTesterConfig()
 	cfg.Countries = append(cfg.Countries, "CA")
@@ -1156,8 +1188,6 @@ func assertStatusCode(t *testing.T, req *http.Response, expected int) {
 func assertRequestHeader(t *testing.T, req *http.Request, key string, expected string) {
 	t.Helper()
 
-	fmt.Println(req.Header.Get(key))
-
 	if received := req.Header.Get(key); received != expected {
 		t.Errorf("header value mismatch: %s: %s <> %s", key, expected, received)
 	}