diff --git a/pkg/coredata/ip_country_block.go b/pkg/coredata/ip_country_block.go new file mode 100644 index 000000000..5512cda98 --- /dev/null +++ b/pkg/coredata/ip_country_block.go @@ -0,0 +1,98 @@ +// Copyright (c) 2026 Probo Inc . +// +// 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 coredata + +import ( + "context" + "fmt" + + "github.com/jackc/pgx/v5" + "go.gearno.de/kit/pg" +) + +type IPCountryBlock struct { + CIDR string `db:"cidr"` + CountryCode CountryCode `db:"country_code"` +} + +func LookupCountryByIP(ctx context.Context, conn pg.Querier, ip string) (CountryCode, error) { + 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[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 IsIPCountryBlocksPopulated(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 TruncateIPCountryBlocks(ctx context.Context, conn pg.Querier) error { + _, err := conn.Exec(ctx, "TRUNCATE common_ip_country_blocks") + if err != nil { + return fmt.Errorf("cannot truncate common_ip_country_blocks: %w", err) + } + return nil +} + +func CopyIPCountryBlocks(ctx context.Context, conn pg.Querier, blocks []IPCountryBlock) error { + rows := make([][]any, len(blocks)) + for i, b := range blocks { + rows[i] = []any{b.CIDR, b.CountryCode.String()} + } + + _, err := conn.CopyFrom( + ctx, + pgx.Identifier{"common_ip_country_blocks"}, + []string{"cidr", "country_code"}, + pgx.CopyFromRows(rows), + ) + if err != nil { + return fmt.Errorf("cannot copy rows into common_ip_country_blocks: %w", err) + } + + return nil +} diff --git a/pkg/geoloc/service.go b/pkg/geoloc/service.go index a4c8ce571..f8e36412d 100644 --- a/pkg/geoloc/service.go +++ b/pkg/geoloc/service.go @@ -23,7 +23,6 @@ import ( "path/filepath" "strings" - "github.com/jackc/pgx/v5" "go.gearno.de/kit/pg" "go.probo.inc/probo/pkg/coredata" ) @@ -44,12 +43,7 @@ func (s *Service) ImportFromDir(ctx context.Context, dataDir string) error { return fmt.Errorf("cannot read country directory: %w", err) } - type row struct { - cidr string - countryCode coredata.CountryCode - } - - var rows []row + var blocks []coredata.IPCountryBlock for _, entry := range entries { if !entry.IsDir() { @@ -71,7 +65,10 @@ func (s *Service) ImportFromDir(ctx context.Context, dataDir string) error { } for _, cidr := range cidrs { - rows = append(rows, row{cidr: cidr, countryCode: cc}) + blocks = append(blocks, coredata.IPCountryBlock{ + CIDR: cidr, + CountryCode: cc, + }) } } } @@ -79,24 +76,12 @@ func (s *Service) ImportFromDir(ctx context.Context, dataDir string) error { 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) + if err := coredata.TruncateIPCountryBlocks(ctx, tx); err != nil { + return 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) + if err := coredata.CopyIPCountryBlocks(ctx, tx, blocks); err != nil { + return err } return nil @@ -110,46 +95,11 @@ func (s *Service) LookupCountry(ctx context.Context, conn pg.Querier, ip string) 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 + return coredata.LookupCountryByIP(ctx, conn, ip) } 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 + return coredata.IsIPCountryBlocksPopulated(ctx, conn) } func parseCIDRFile(path string) ([]string, error) {