add geoip country/asn labels and ipv6

internal/dnstt/geoip.go

@@ -0,0 +1,108 @@
+package dnstt
+
+import (
+	"fmt"
+	"net/netip"
+	"strconv"
+
+	geoip2 "github.com/oschwald/geoip2-golang/v2"
+)
+
+const (
+	// UnknownCountry is used when country lookup is enabled but no country is found.
+	UnknownCountry = "ZZ"
+	// UnknownASN is used when ASN lookup is enabled but no ASN is found.
+	UnknownASN = "0"
+)
+
+// GeoLabels are optional labels derived from a resolver IP address.
+type GeoLabels struct {
+	Country string
+	ASN     string
+}
+
+// GeoResolver resolves optional GeoIP labels for resolver IP addresses.
+type GeoResolver interface {
+	Lookup(netip.Addr) GeoLabels
+	LabelNames() []string
+}
+
+// GeoIPResolver uses optional MaxMind Country and ASN databases.
+type GeoIPResolver struct {
+	countryDB *geoip2.Reader
+	asnDB     *geoip2.Reader
+}
+
+// OpenGeoIPResolver opens optional MaxMind Country and ASN databases.
+func OpenGeoIPResolver(countryDatabase string, asnDatabase string) (*GeoIPResolver, error) {
+	resolver := &GeoIPResolver{}
+	if countryDatabase != "" {
+		db, err := geoip2.Open(countryDatabase)
+		if err != nil {
+			return nil, fmt.Errorf("open country database: %w", err)
+		}
+		resolver.countryDB = db
+	}
+	if asnDatabase != "" {
+		db, err := geoip2.Open(asnDatabase)
+		if err != nil {
+			resolver.Close()
+			return nil, fmt.Errorf("open ASN database: %w", err)
+		}
+		resolver.asnDB = db
+	}
+	return resolver, nil
+}
+
+// Close closes any open MaxMind databases.
+func (r *GeoIPResolver) Close() {
+	if r.countryDB != nil {
+		r.countryDB.Close()
+	}
+	if r.asnDB != nil {
+		r.asnDB.Close()
+	}
+}
+
+// LabelNames returns the optional Prometheus labels enabled by configured databases.
+func (r *GeoIPResolver) LabelNames() []string {
+	var labels []string
+	if r.countryDB != nil {
+		labels = append(labels, "country")
+	}
+	if r.asnDB != nil {
+		labels = append(labels, "asn")
+	}
+	return labels
+}
+
+// Lookup returns optional GeoIP labels for a resolver IP address.
+func (r *GeoIPResolver) Lookup(addr netip.Addr) GeoLabels {
+	labels := GeoLabels{}
+	if !addr.IsValid() {
+		if r.countryDB != nil {
+			labels.Country = UnknownCountry
+		}
+		if r.asnDB != nil {
+			labels.ASN = UnknownASN
+		}
+		return labels
+	}
+
+	if r.countryDB != nil {
+		labels.Country = UnknownCountry
+		record, err := r.countryDB.Country(addr)
+		if err == nil && record.HasData() && record.Country.ISOCode != "" {
+			labels.Country = record.Country.ISOCode
+		}
+	}
+	if r.asnDB != nil {
+		labels.ASN = UnknownASN
+		record, err := r.asnDB.ASN(addr)
+		if err == nil && record.HasData() && record.AutonomousSystemNumber != 0 {
+			labels.ASN = strconv.FormatUint(uint64(record.AutonomousSystemNumber), 10)
+		}
+	}
+
+	return labels
+}