151 lines
4.4 KiB
Go
151 lines
4.4 KiB
Go
// Copyright (c) 2026 Probo Inc <hello@probo.com>.
|
|
//
|
|
// Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
// of this software and associated documentation files (the "Software"), to deal
|
|
// in the Software without restriction, including without limitation the rights
|
|
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
// copies of the Software, and to permit persons to whom the Software is
|
|
// furnished to do so, subject to the following conditions:
|
|
//
|
|
// The above copyright notice and this permission notice shall be included in
|
|
// all copies or substantial portions of the Software.
|
|
//
|
|
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
// 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
|
|
}
|
|
|
|
const ipCountryBlocksStagingTable = "common_ip_country_blocks_staging"
|
|
|
|
func CreateIPCountryBlocksStaging(ctx context.Context, conn pg.Querier) error {
|
|
q := `
|
|
DROP TABLE IF EXISTS common_ip_country_blocks_staging;
|
|
|
|
ALTER INDEX IF EXISTS idx_common_ip_country_blocks_staging_cidr
|
|
RENAME TO idx_common_ip_country_blocks_cidr;
|
|
|
|
CREATE TABLE common_ip_country_blocks_staging (LIKE common_ip_country_blocks INCLUDING DEFAULTS);
|
|
`
|
|
|
|
if _, err := conn.Exec(ctx, q); err != nil {
|
|
return fmt.Errorf("cannot create staging table: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func CopyIPCountryBlocksStaging(ctx context.Context, conn pg.Querier, blocks []IPCountryBlock) error {
|
|
_, err := conn.CopyFrom(
|
|
ctx,
|
|
pgx.Identifier{ipCountryBlocksStagingTable},
|
|
[]string{"cidr", "country_code"},
|
|
pgx.CopyFromSlice(len(blocks), func(i int) ([]any, error) {
|
|
return []any{blocks[i].CIDR, blocks[i].CountryCode.String()}, nil
|
|
}),
|
|
)
|
|
if err != nil {
|
|
return fmt.Errorf("cannot copy rows into staging table: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func FinalizeIPCountryBlocksStaging(ctx context.Context, conn pg.Querier) error {
|
|
q := `
|
|
CREATE INDEX idx_common_ip_country_blocks_staging_cidr
|
|
ON common_ip_country_blocks_staging USING gist (cidr inet_ops);
|
|
ANALYZE common_ip_country_blocks_staging;
|
|
`
|
|
|
|
if _, err := conn.Exec(ctx, q); err != nil {
|
|
return fmt.Errorf("cannot finalize staging table: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func SwapIPCountryBlocksStaging(ctx context.Context, conn pg.Querier) error {
|
|
q := `
|
|
DROP TABLE common_ip_country_blocks;
|
|
ALTER TABLE common_ip_country_blocks_staging RENAME TO common_ip_country_blocks;
|
|
ALTER INDEX idx_common_ip_country_blocks_staging_cidr
|
|
RENAME TO idx_common_ip_country_blocks_cidr;
|
|
`
|
|
|
|
if _, err := conn.Exec(ctx, q); err != nil {
|
|
return fmt.Errorf("cannot swap staging table: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func DropIPCountryBlocksStaging(ctx context.Context, conn pg.Querier) error {
|
|
if _, err := conn.Exec(ctx, "DROP TABLE IF EXISTS common_ip_country_blocks_staging"); err != nil {
|
|
return fmt.Errorf("cannot drop staging table: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|