feat(thoth): add GeoIP checker

Signed-off-by: Xe Iaso <me@xeiaso.net>
This commit is contained in:
Xe Iaso
2025-05-22 10:48:15 -04:00
parent 4aea22fac5
commit 49ab76c9dd
2 changed files with 99 additions and 0 deletions

View File

@@ -0,0 +1,36 @@
package thoth
import (
"context"
"net/http"
"strings"
"time"
iptoasnv1 "github.com/TecharoHQ/thoth-proto/gen/techaro/thoth/iptoasn/v1"
)
type GeoIPChecker struct {
iptoasn iptoasnv1.IpToASNServiceClient
countries map[string]struct{}
hash string
}
func (gipc *GeoIPChecker) Check(r *http.Request) (bool, error) {
ctx, cancel := context.WithTimeout(r.Context(), 50*time.Millisecond)
defer cancel()
ipInfo, err := gipc.iptoasn.Lookup(ctx, &iptoasnv1.LookupRequest{
IpAddress: r.Header.Get("X-Real-Ip"),
})
if err != nil {
return false, err
}
_, ok := gipc.countries[strings.ToLower(ipInfo.GetCountryCode())]
return ok, nil
}
func (gipc *GeoIPChecker) Hash() string {
return gipc.hash
}

View File

@@ -0,0 +1,63 @@
package thoth
import (
"fmt"
"net/http/httptest"
"testing"
"github.com/TecharoHQ/anubis/lib/policy"
)
var _ policy.Checker = &ASNChecker{}
func TestGeoIPChecker(t *testing.T) {
cli := loadSecrets(t)
asnc := &GeoIPChecker{
iptoasn: cli.iptoasn,
countries: map[string]struct{}{
"us": {},
},
hash: "foobar",
}
for _, cs := range []struct {
ipAddress string
wantMatch bool
wantError bool
}{
{
ipAddress: "1.1.1.1",
wantMatch: true,
wantError: false,
},
{
ipAddress: "70.31.0.1",
wantMatch: false,
wantError: false,
},
{
ipAddress: "taco",
wantMatch: false,
wantError: true,
},
} {
t.Run(fmt.Sprintf("%v", cs), func(t *testing.T) {
req := httptest.NewRequest("GET", "/", nil)
req.Header.Set("X-Real-Ip", cs.ipAddress)
match, err := asnc.Check(req)
if match != cs.wantMatch {
t.Errorf("Wanted match: %v, got: %v", cs.wantMatch, match)
}
switch {
case err != nil && !cs.wantError:
t.Errorf("Did not want error but got: %v", err)
case err == nil && cs.wantError:
t.Error("Wanted error but got none")
}
})
}
}