Add IP-to-country geolocation service
Introduce a geoloc package that stores CIDR-to-country mappings in PostgreSQL using the native cidr type with a GiST index for fast containment lookups. Data comes from the ipverse/country-ip-blocks dataset added as a git submodule. A standalone geoloc-import command reads the TXT files from disk and bulk-loads them via COPY. probod wires the service and logs a warning when the table is empty. Signed-off-by: Émile Ré <emile@getprobo.com>
This commit is contained in:
3
.gitmodules
vendored
3
.gitmodules
vendored
@@ -1,3 +1,6 @@
|
||||
[submodule "pkg/validator/data/disposable-email-domains"]
|
||||
path = pkg/validator/data/disposable-email-domains
|
||||
url = https://github.com/disposable-email-domains/disposable-email-domains.git
|
||||
[submodule "pkg/geoloc/data/country-ip-blocks"]
|
||||
path = pkg/geoloc/data/country-ip-blocks
|
||||
url = https://github.com/ipverse/country-ip-blocks.git
|
||||
|
||||
104
cmd/geoloc-import/main.go
Normal file
104
cmd/geoloc-import/main.go
Normal file
@@ -0,0 +1,104 @@
|
||||
// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
|
||||
//
|
||||
// Permission to use, copy, modify, and/or distribute this software for any
|
||||
// purpose with or without fee is hereby granted, provided that the above
|
||||
// copyright notice and this permission notice appear in all copies.
|
||||
//
|
||||
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
|
||||
// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
|
||||
// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
|
||||
// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
|
||||
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
|
||||
// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
|
||||
// PERFORMANCE OF THIS SOFTWARE.
|
||||
|
||||
// Command geoloc-import loads IP-to-country CIDR blocks from the
|
||||
// ipverse/country-ip-blocks dataset into the common_ip_country_blocks table.
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"flag"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"os"
|
||||
|
||||
"go.gearno.de/kit/pg"
|
||||
"go.probo.inc/probo/pkg/geoloc"
|
||||
)
|
||||
|
||||
func main() {
|
||||
if err := run(); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "error: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func run() error {
|
||||
var (
|
||||
pgDSN string
|
||||
dataDir string
|
||||
)
|
||||
|
||||
flag.StringVar(
|
||||
&pgDSN,
|
||||
"pg-dsn",
|
||||
os.Getenv("DATABASE_URL"),
|
||||
"PostgreSQL connection URL (default: DATABASE_URL env)",
|
||||
)
|
||||
flag.StringVar(
|
||||
&dataDir,
|
||||
"data-dir",
|
||||
"pkg/geoloc/data/country-ip-blocks",
|
||||
"path to ipverse country-ip-blocks checkout",
|
||||
)
|
||||
flag.Parse()
|
||||
|
||||
if pgDSN == "" {
|
||||
return fmt.Errorf("set -pg-dsn or DATABASE_URL")
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
pgClient, err := newPgClientFromDSN(pgDSN)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot create pg client: %w", err)
|
||||
}
|
||||
|
||||
svc := geoloc.NewService(pgClient)
|
||||
|
||||
fmt.Printf("importing IP country blocks from %s\n", dataDir)
|
||||
|
||||
if err := svc.ImportFromDir(ctx, dataDir); err != nil {
|
||||
return fmt.Errorf("cannot import geoloc data: %w", err)
|
||||
}
|
||||
|
||||
fmt.Println("done")
|
||||
return nil
|
||||
}
|
||||
|
||||
func newPgClientFromDSN(dsn string) (*pg.Client, error) {
|
||||
u, err := url.Parse(dsn)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot parse DSN: %w", err)
|
||||
}
|
||||
|
||||
var opts []pg.Option
|
||||
|
||||
if u.Host != "" {
|
||||
opts = append(opts, pg.WithAddr(u.Host))
|
||||
}
|
||||
|
||||
if u.User != nil {
|
||||
opts = append(opts, pg.WithUser(u.User.Username()))
|
||||
if password, ok := u.User.Password(); ok {
|
||||
opts = append(opts, pg.WithPassword(password))
|
||||
}
|
||||
}
|
||||
|
||||
if len(u.Path) > 1 {
|
||||
opts = append(opts, pg.WithDatabase(u.Path[1:]))
|
||||
}
|
||||
|
||||
return pg.NewClient(opts...)
|
||||
}
|
||||
21
pkg/coredata/migrations/20260506T084741Z.sql
Normal file
21
pkg/coredata/migrations/20260506T084741Z.sql
Normal file
@@ -0,0 +1,21 @@
|
||||
-- Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
|
||||
--
|
||||
-- Permission to use, copy, modify, and/or distribute this software for any
|
||||
-- purpose with or without fee is hereby granted, provided that the above
|
||||
-- copyright notice and this permission notice appear in all copies.
|
||||
--
|
||||
-- THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
|
||||
-- REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
|
||||
-- AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
|
||||
-- INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
|
||||
-- LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
|
||||
-- OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
|
||||
-- PERFORMANCE OF THIS SOFTWARE.
|
||||
|
||||
CREATE TABLE common_ip_country_blocks (
|
||||
cidr CIDR NOT NULL,
|
||||
country_code CHAR(2) NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX idx_common_ip_country_blocks_cidr
|
||||
ON common_ip_country_blocks USING gist (cidr inet_ops);
|
||||
1
pkg/geoloc/data/country-ip-blocks
Submodule
1
pkg/geoloc/data/country-ip-blocks
Submodule
Submodule pkg/geoloc/data/country-ip-blocks added at b95fe72214
184
pkg/geoloc/service.go
Normal file
184
pkg/geoloc/service.go
Normal file
@@ -0,0 +1,184 @@
|
||||
// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
|
||||
//
|
||||
// Permission to use, copy, modify, and/or distribute this software for any
|
||||
// purpose with or without fee is hereby granted, provided that the above
|
||||
// copyright notice and this permission notice appear in all copies.
|
||||
//
|
||||
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
|
||||
// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
|
||||
// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
|
||||
// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
|
||||
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
|
||||
// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
|
||||
// PERFORMANCE OF THIS SOFTWARE.
|
||||
|
||||
package geoloc
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"go.gearno.de/kit/pg"
|
||||
"go.probo.inc/probo/pkg/coredata"
|
||||
)
|
||||
|
||||
type Service struct {
|
||||
pgClient *pg.Client
|
||||
}
|
||||
|
||||
func NewService(pgClient *pg.Client) *Service {
|
||||
return &Service{pgClient: pgClient}
|
||||
}
|
||||
|
||||
func (s *Service) ImportFromDir(ctx context.Context, dataDir string) error {
|
||||
countryDir := filepath.Join(dataDir, "country")
|
||||
|
||||
entries, err := os.ReadDir(countryDir)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot read country directory: %w", err)
|
||||
}
|
||||
|
||||
type row struct {
|
||||
cidr string
|
||||
countryCode coredata.CountryCode
|
||||
}
|
||||
|
||||
var rows []row
|
||||
|
||||
for _, entry := range entries {
|
||||
if !entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
|
||||
code := strings.ToUpper(entry.Name())
|
||||
var cc coredata.CountryCode
|
||||
if err := cc.Scan(code); err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
for _, filename := range []string{"ipv4-aggregated.txt", "ipv6-aggregated.txt"} {
|
||||
path := filepath.Join(countryDir, entry.Name(), filename)
|
||||
|
||||
cidrs, err := parseCIDRFile(path)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
for _, cidr := range cidrs {
|
||||
rows = append(rows, row{cidr: cidr, countryCode: cc})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return s.pgClient.WithTx(
|
||||
ctx,
|
||||
func(ctx context.Context, tx pg.Tx) error {
|
||||
_, err := tx.Exec(ctx, "TRUNCATE common_ip_country_blocks")
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot truncate common_ip_country_blocks: %w", err)
|
||||
}
|
||||
|
||||
pgxRows := make([][]any, len(rows))
|
||||
for i, r := range rows {
|
||||
pgxRows[i] = []any{r.cidr, r.countryCode.String()}
|
||||
}
|
||||
|
||||
_, err = tx.CopyFrom(
|
||||
ctx,
|
||||
pgx.Identifier{"common_ip_country_blocks"},
|
||||
[]string{"cidr", "country_code"},
|
||||
pgx.CopyFromRows(pgxRows),
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot copy rows into common_ip_country_blocks: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
func (s *Service) LookupCountry(ctx context.Context, conn pg.Querier, ip string) (coredata.CountryCode, error) {
|
||||
parsed := net.ParseIP(ip)
|
||||
if parsed == nil {
|
||||
return "", fmt.Errorf("cannot parse IP address: %q", ip)
|
||||
}
|
||||
|
||||
q := `
|
||||
SELECT country_code
|
||||
FROM common_ip_country_blocks
|
||||
WHERE cidr >>= @ip::inet
|
||||
ORDER BY masklen(cidr) DESC
|
||||
LIMIT 1;
|
||||
`
|
||||
|
||||
args := pgx.StrictNamedArgs{"ip": ip}
|
||||
|
||||
rows, err := conn.Query(ctx, q, args)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot query ip country blocks: %w", err)
|
||||
}
|
||||
|
||||
cc, err := pgx.CollectExactlyOneRow(rows, pgx.RowTo[coredata.CountryCode])
|
||||
if err != nil {
|
||||
if err == pgx.ErrNoRows {
|
||||
return "", nil
|
||||
}
|
||||
return "", fmt.Errorf("cannot collect ip country block row: %w", err)
|
||||
}
|
||||
|
||||
return cc, nil
|
||||
}
|
||||
|
||||
func (s *Service) IsPopulated(ctx context.Context, conn pg.Querier) (bool, error) {
|
||||
q := `SELECT EXISTS (SELECT 1 FROM common_ip_country_blocks);`
|
||||
|
||||
rows, err := conn.Query(ctx, q)
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("cannot check if ip country blocks is populated: %w", err)
|
||||
}
|
||||
|
||||
populated, err := pgx.CollectExactlyOneRow(rows, pgx.RowTo[bool])
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("cannot collect populated check: %w", err)
|
||||
}
|
||||
|
||||
return populated, nil
|
||||
}
|
||||
|
||||
func parseCIDRFile(path string) ([]string, error) {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
var cidrs []string
|
||||
scanner := bufio.NewScanner(f)
|
||||
|
||||
for scanner.Scan() {
|
||||
line := strings.TrimSpace(scanner.Text())
|
||||
if line == "" || strings.HasPrefix(line, "#") {
|
||||
continue
|
||||
}
|
||||
|
||||
_, _, err := net.ParseCIDR(line)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
cidrs = append(cidrs, line)
|
||||
}
|
||||
|
||||
if err := scanner.Err(); err != nil {
|
||||
return nil, fmt.Errorf("cannot scan file: %w", err)
|
||||
}
|
||||
|
||||
return cidrs, nil
|
||||
}
|
||||
@@ -57,6 +57,7 @@ import (
|
||||
"go.probo.inc/probo/pkg/evidencedescriber"
|
||||
"go.probo.inc/probo/pkg/file"
|
||||
"go.probo.inc/probo/pkg/filemanager"
|
||||
"go.probo.inc/probo/pkg/geoloc"
|
||||
"go.probo.inc/probo/pkg/html2pdf"
|
||||
"go.probo.inc/probo/pkg/iam"
|
||||
"go.probo.inc/probo/pkg/iam/oauth2server"
|
||||
@@ -260,6 +261,24 @@ func (impl *Implm) Run(
|
||||
return fmt.Errorf("cannot migrate database schema: %w", err)
|
||||
}
|
||||
|
||||
geolocService := geoloc.NewService(pgClient)
|
||||
err = pgClient.WithConn(
|
||||
ctx,
|
||||
func(ctx context.Context, conn pg.Querier) error {
|
||||
populated, err := geolocService.IsPopulated(ctx, conn)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !populated {
|
||||
l.Warn("IP geolocation table is empty; run geoloc-import to populate it")
|
||||
}
|
||||
return nil
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
l.ErrorCtx(ctx, "cannot check geoloc table", log.Error(err))
|
||||
}
|
||||
|
||||
hp, err := passwdhash.NewProfile(pepper, uint32(impl.cfg.Auth.Password.Iterations))
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot create hashing profile: %w", err)
|
||||
|
||||
Reference in New Issue
Block a user