84 lines
2.1 KiB
Go
84 lines
2.1 KiB
Go
// Package lookuptest writes lookup databases, IPinfo Lite files in their
|
|
// .mmdb form, for the tests of the packages that read them.
|
|
package lookuptest
|
|
|
|
import (
|
|
"bytes"
|
|
"net"
|
|
"os"
|
|
"testing"
|
|
|
|
"github.com/maxmind/mmdbwriter"
|
|
"github.com/maxmind/mmdbwriter/mmdbtype"
|
|
)
|
|
|
|
// fileMode is the mode of the files written: read and written by their
|
|
// owner alone.
|
|
const fileMode = 0o600
|
|
|
|
// Network is what a lookup database holds about a netblock, of the fields
|
|
// sww reads: its AS number, such as AS64496, the AS's name, and
|
|
// its country, such as DE.
|
|
type Network struct {
|
|
ASN string
|
|
ASName string
|
|
Country string
|
|
}
|
|
|
|
// Write writes a lookup database at path that places each netblock in
|
|
// networks, such as 203.0.113.0/24, as its Network says, and no other
|
|
// address.
|
|
func Write(tb testing.TB, path string, networks map[string]Network) {
|
|
tb.Helper()
|
|
|
|
records := make(map[string]mmdbtype.Map, len(networks))
|
|
for netblock, network := range networks {
|
|
records[netblock] = mmdbtype.Map{
|
|
"asn": mmdbtype.String(network.ASN),
|
|
"as_name": mmdbtype.String(network.ASName),
|
|
"country_code": mmdbtype.String(network.Country),
|
|
}
|
|
}
|
|
|
|
WriteRecords(tb, path, records)
|
|
}
|
|
|
|
// WriteRecords writes a lookup database at path that holds each record in
|
|
// records for its netblock, and nothing for any other address.
|
|
func WriteRecords(tb testing.TB, path string, records map[string]mmdbtype.Map) {
|
|
tb.Helper()
|
|
|
|
tree, err := mmdbwriter.New(mmdbwriter.Options{
|
|
DatabaseType: "ipinfo_lite",
|
|
// The tests' clients are in the netblocks kept for documentation.
|
|
IncludeReservedNetworks: true,
|
|
})
|
|
if err != nil {
|
|
tb.Fatalf("new lookup database: %v", err)
|
|
}
|
|
|
|
for netblock, record := range records {
|
|
_, network, err := net.ParseCIDR(netblock)
|
|
if err != nil {
|
|
tb.Fatalf("netblock %q: %v", netblock, err)
|
|
}
|
|
|
|
err = tree.Insert(network, record)
|
|
if err != nil {
|
|
tb.Fatalf("insert %s: %v", netblock, err)
|
|
}
|
|
}
|
|
|
|
var database bytes.Buffer
|
|
|
|
_, err = tree.WriteTo(&database)
|
|
if err != nil {
|
|
tb.Fatalf("write the lookup database: %v", err)
|
|
}
|
|
|
|
err = os.WriteFile(path, database.Bytes(), fileMode)
|
|
if err != nil {
|
|
tb.Fatalf("write %s: %v", path, err)
|
|
}
|
|
}
|