asn lookup support

This commit is contained in:
2025-07-31 19:18:32 +02:00
parent c269dd0f82
commit dfec9ba602
3 changed files with 200 additions and 47 deletions
+107 -13
View File
@@ -5,6 +5,7 @@ import (
"log"
"net"
"net/http"
"sort"
"strings"
"sync"
"time"
@@ -24,14 +25,40 @@ type DataStruct struct {
Loc *string `json:"loc"`
}
// Global IP cache with 10 minute TTL
type ASNDataResponse struct {
ASNDetails ASNDetails `json:"asn_details"`
Prefixes PrefixInfo `json:"prefixes"`
SourceDetails SourceDetails `json:"source_details"`
}
type ASNDetails struct {
ASN uint `json:"asn"`
Name string `json:"name"`
}
type PrefixInfo struct {
IPv4 []string `json:"ipv4"`
IPv6 []string `json:"ipv6"`
}
type SourceDetails struct {
Source string `json:"source"`
}
// Global caches with 10 minute TTL
var ipCache = NewIPCache(10 * time.Minute)
var asnCache = NewASNCache(10 * time.Minute)
type cachedIPData struct {
data *DataStruct
time time.Time
}
type cachedASNData struct {
data *ASNDataResponse
time time.Time
}
// IPCache provides thread-safe caching of IP lookup results
type IPCache struct {
cache sync.Map
@@ -45,7 +72,6 @@ func NewIPCache(ttl time.Duration) *IPCache {
}
}
// Set stores an IP in the cache
func (c *IPCache) Set(ipStr string, data *DataStruct) {
c.cache.Store(ipStr, cachedIPData{
data: data,
@@ -53,7 +79,6 @@ func (c *IPCache) Set(ipStr string, data *DataStruct) {
})
}
// Get retrieves an IP from the cache if it exists and is not expired
func (c *IPCache) Get(ipStr string) (*DataStruct, bool) {
if cachedData, ok := c.cache.Load(ipStr); ok {
cached := cachedData.(cachedIPData)
@@ -65,9 +90,37 @@ func (c *IPCache) Get(ipStr string) (*DataStruct, bool) {
return nil, false
}
type ASNCache struct {
cache sync.Map
ttl time.Duration
}
func NewASNCache(ttl time.Duration) *ASNCache {
return &ASNCache{
ttl: ttl,
}
}
func (c *ASNCache) Set(asn uint, data *ASNDataResponse) {
c.cache.Store(asn, cachedASNData{
data: data,
time: time.Now(),
})
}
func (c *ASNCache) Get(asn uint) (*ASNDataResponse, bool) {
if cachedData, ok := c.cache.Load(asn); ok {
cached := cachedData.(cachedASNData)
if time.Since(cached.time) < c.ttl {
return cached.data, true
}
c.cache.Delete(asn)
}
return nil, false
}
// LookupIPData looks up IP data in the databases with caching
func LookupIPData(geoIP *db.GeoIPManager, ip net.IP) *DataStruct {
// Check cache first
if data, found := ipCache.Get(ip.String()); found {
return data
}
@@ -90,7 +143,6 @@ func LookupIPData(geoIP *db.GeoIPManager, ip net.IP) *DataStruct {
} `maxminddb:"location"`
}
// Get database readers using thread-safe accessor methods
cityDB := geoIP.GetCityDB()
err := cityDB.Lookup(ip, &cityRecord)
if err != nil {
@@ -98,10 +150,7 @@ func LookupIPData(geoIP *db.GeoIPManager, ip net.IP) *DataStruct {
return nil
}
var asnRecord struct {
AutonomousSystemNumber uint `maxminddb:"autonomous_system_number"`
AutonomousSystemOrganization string `maxminddb:"autonomous_system_organization"`
}
var asnRecord db.ASNRecord
asnDB := geoIP.GetASNDB()
err = asnDB.Lookup(ip, &asnRecord)
if err != nil {
@@ -130,12 +179,59 @@ func LookupIPData(geoIP *db.GeoIPManager, ip net.IP) *DataStruct {
Loc: ToPtr(fmt.Sprintf("%.4f,%.4f", cityRecord.Location.Latitude, cityRecord.Location.Longitude)),
}
// Store in cache
ipCache.Set(ip.String(), data)
return data
}
func LookupASNData(geoIP *db.GeoIPManager, targetASN uint) (*ASNDataResponse, error) {
if data, found := asnCache.Get(targetASN); found {
return data, nil
}
prefixes := geoIP.GetASNPrefixes(targetASN)
if len(prefixes) == 0 {
return nil, fmt.Errorf("no prefixes found for ASN %d in the database", targetASN)
}
var orgName string
var ipv4Prefixes, ipv6Prefixes []string
var record db.ASNRecord
if err := geoIP.GetASNDB().Lookup(prefixes[0].IP, &record); err == nil {
orgName = record.AutonomousSystemOrganization
}
for _, prefix := range prefixes {
prefixStr := prefix.String()
if strings.Contains(prefixStr, ":") {
ipv6Prefixes = append(ipv6Prefixes, prefixStr)
} else {
ipv4Prefixes = append(ipv4Prefixes, prefixStr)
}
}
sort.Strings(ipv4Prefixes)
sort.Strings(ipv6Prefixes)
response := &ASNDataResponse{
ASNDetails: ASNDetails{
ASN: targetASN,
Name: orgName,
},
Prefixes: PrefixInfo{
IPv4: ipv4Prefixes,
IPv6: ipv6Prefixes,
},
SourceDetails: SourceDetails{
Source: "GeoLite2-ASN.mmdb",
},
}
asnCache.Set(targetASN, response)
return response, nil
}
// ToPtr converts string to pointer
func ToPtr(s string) *string {
if s == "" {
@@ -156,14 +252,12 @@ func IsBogon(ip net.IP) bool {
// GetRealIP extracts the client's real IP address from request headers
func GetRealIP(r *http.Request) string {
// Try common proxy headers first
for _, header := range []string{"CF-Connecting-IP", "X-Real-IP", "X-Forwarded-For"} {
if ip := r.Header.Get(header); ip != "" {
return strings.TrimSpace(strings.Split(ip, ",")[0])
}
}
// Fall back to remote address
host, _, err := net.SplitHostPort(r.RemoteAddr)
if err != nil {
return r.RemoteAddr
+43 -3
View File
@@ -5,6 +5,7 @@ import (
"errors"
"fmt"
"log"
"net"
"os"
"os/exec"
"sync"
@@ -26,11 +27,17 @@ var (
ErrDownloadFailed = errors.New("failed to download database")
)
type ASNRecord struct {
AutonomousSystemNumber uint `maxminddb:"autonomous_system_number"`
AutonomousSystemOrganization string `maxminddb:"autonomous_system_organization"`
}
// Handles MaxMind GeoIP database operations
type GeoIPManager struct {
cityDB *maxminddb.Reader
asnDB *maxminddb.Reader
mu sync.RWMutex
cityDB *maxminddb.Reader
asnDB *maxminddb.Reader
asnPrefixMap map[uint][]*net.IPNet
mu sync.RWMutex
}
// Creates and initializes a new GeoIP database
@@ -55,9 +62,32 @@ func (g *GeoIPManager) Initialize() error {
return err
}
g.buildASNPrefixMap()
return nil
}
// buildASNPrefixMap iterates the ASN database once and builds the lookup map.
func (g *GeoIPManager) buildASNPrefixMap() {
log.Println("Building ASN prefix map for fast lookups...")
startTime := time.Now()
g.asnPrefixMap = make(map[uint][]*net.IPNet)
networks := g.asnDB.Networks()
for networks.Next() {
var record ASNRecord
subnet, err := networks.Network(&record)
if err != nil {
continue // Skip records that fail to decode
}
if record.AutonomousSystemNumber > 0 {
g.asnPrefixMap[record.AutonomousSystemNumber] = append(g.asnPrefixMap[record.AutonomousSystemNumber], subnet)
}
}
log.Printf("Finished building ASN prefix map in %v", time.Since(startTime))
}
// Opens the city database, downloading it if necessary
func (g *GeoIPManager) openCityDB() error {
var err error
@@ -152,6 +182,13 @@ func (g *GeoIPManager) GetASNDB() *maxminddb.Reader {
return g.asnDB
}
// GetASNPrefixes returns the pre-computed list of prefixes for an ASN.
func (g *GeoIPManager) GetASNPrefixes(asn uint) []*net.IPNet {
g.mu.RLock()
defer g.mu.RUnlock()
return g.asnPrefixMap[asn]
}
// Sets up automatic database updates
func (g *GeoIPManager) StartUpdater(ctx context.Context, updateInterval time.Duration) {
log.Printf("Starting MaxMind GeoIP database updater with interval: %s", updateInterval)
@@ -209,6 +246,9 @@ func (g *GeoIPManager) UpdateDatabases() error {
return fmt.Errorf("reopening ASN database: %w", err)
}
// Rebuild the prefix map with the new data.
g.buildASNPrefixMap()
log.Println("Successfully updated and reloaded GeoIP databases")
return nil
}
+50 -31
View File
@@ -9,6 +9,7 @@ import (
"net/http"
"os"
"os/signal"
"strconv"
"strings"
"syscall"
"time"
@@ -32,14 +33,12 @@ func (w gzipResponseWriter) Write(b []byte) (int, error) {
return w.Writer.Write(b)
}
// HTTP server and its dependencies
type Server struct {
geoIP *db.GeoIPManager
server *http.Server
shutdownSignal chan struct{}
}
// Creates a new server with the given GeoIPManager
func NewServer(geoIP *db.GeoIPManager) *Server {
s := &Server{
geoIP: geoIP,
@@ -47,10 +46,9 @@ func NewServer(geoIP *db.GeoIPManager) *Server {
}
mux := http.NewServeMux()
mux.Handle("/health", utils.HealthCheck()) // Register healthcheck
mux.HandleFunc("/", s.handler) // Main handler
mux.Handle("/health", utils.HealthCheck())
mux.HandleFunc("/", s.router)
// Create HTTP server
s.server = &http.Server{
Addr: ":3000",
Handler: mux,
@@ -62,9 +60,7 @@ func NewServer(geoIP *db.GeoIPManager) *Server {
return s
}
// Starts the HTTP server
func (s *Server) Start(ctx context.Context) error {
// Start the server in a goroutine
go func() {
log.Println("Server listening on :3000")
if err := s.server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
@@ -72,7 +68,6 @@ func (s *Server) Start(ctx context.Context) error {
}
}()
// Wait for shutdown signal or context cancellation
select {
case <-s.shutdownSignal:
log.Println("Shutdown requested internally")
@@ -80,11 +75,9 @@ func (s *Server) Start(ctx context.Context) error {
log.Println("Shutdown requested from context")
}
// Create shutdown context with timeout
shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
// Shutdown the server
if err := s.server.Shutdown(shutdownCtx); err != nil {
return err
}
@@ -93,12 +86,10 @@ func (s *Server) Start(ctx context.Context) error {
return nil
}
// Graceful server shutdown
func (s *Server) Shutdown() {
close(s.shutdownSignal)
}
// Field access functions map
var fieldMap = map[string]func(*common.DataStruct) *string{
"ip": func(d *common.DataStruct) *string { return d.IP },
"hostname": func(d *common.DataStruct) *string { return d.Hostname },
@@ -110,7 +101,6 @@ var fieldMap = map[string]func(*common.DataStruct) *string{
"loc": func(d *common.DataStruct) *string { return d.Loc },
}
// Retrieves a field from the dataStruct using the fieldMap
func getField(data *common.DataStruct, field string) *string {
if f, ok := fieldMap[field]; ok {
return f(data)
@@ -118,11 +108,7 @@ func getField(data *common.DataStruct, field string) *string {
return nil
}
// Processes HTTP requests
func (s *Server) handler(w http.ResponseWriter, r *http.Request) {
requestedThings := strings.Split(strings.Trim(r.URL.Path, "/"), "/")
// Enable gzip compression if requested
func (s *Server) router(w http.ResponseWriter, r *http.Request) {
if strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
w.Header().Set("Content-Encoding", "gzip")
gz := gzip.NewWriter(w)
@@ -130,9 +116,54 @@ func (s *Server) handler(w http.ResponseWriter, r *http.Request) {
w = &gzipResponseWriter{Writer: gz, ResponseWriter: w}
}
path := strings.Trim(r.URL.Path, "/")
lowerPath := strings.ToLower(path)
if strings.HasPrefix(lowerPath, "as") {
s.handleASNLookup(w, r, path)
return
}
s.handleIPLookup(w, r, path)
}
func (s *Server) handleASNLookup(w http.ResponseWriter, _ *http.Request, path string) {
var asnStr string
lowerPath := strings.ToLower(path)
if strings.HasPrefix(lowerPath, "asn/") {
asnStr = path[4:]
} else if strings.HasPrefix(lowerPath, "as") {
asnStr = path[2:]
} else {
sendJSONError(w, "Invalid ASN query format. Use /asn/<number> or /AS<number>.", http.StatusBadRequest)
return
}
asn, err := strconv.ParseUint(asnStr, 10, 32)
if err != nil || asn == 0 {
sendJSONError(w, "Invalid ASN: must be a positive number.", http.StatusBadRequest)
return
}
data, err := common.LookupASNData(s.geoIP, uint(asn))
if err != nil {
if strings.Contains(err.Error(), "no prefixes found") {
sendJSONError(w, err.Error(), http.StatusNotFound)
} else {
log.Printf("Error looking up ASN data for %d: %v", asn, err)
sendJSONError(w, "Error retrieving data for ASN.", http.StatusInternalServerError)
}
return
}
sendJSONResponse(w, data, http.StatusOK)
}
func (s *Server) handleIPLookup(w http.ResponseWriter, r *http.Request, path string) {
requestedThings := strings.Split(path, "/")
var IPAddress, field string
// Parse the request URL
switch len(requestedThings) {
case 0:
IPAddress = common.GetRealIP(r) // Default to visitor's IP
@@ -161,38 +192,32 @@ func (s *Server) handler(w http.ResponseWriter, r *http.Request) {
return
}
// Validate the resolved IP
ip := net.ParseIP(IPAddress)
if ip == nil {
sendJSONError(w, "Please provide a valid IP address.", http.StatusBadRequest)
return
}
// Check if the IP is bogon
if common.IsBogon(ip) {
sendJSONResponse(w, bogonDataStruct{IP: ip.String(), Bogon: true}, http.StatusOK)
return
}
// Look up IP data
data := common.LookupIPData(s.geoIP, ip)
if data == nil {
sendJSONError(w, "Please provide a valid IP address.", http.StatusBadRequest)
return
}
// Handle specific field requests
if field != "" {
value := getField(data, field)
sendJSONResponse(w, map[string]*string{field: value}, http.StatusOK)
return
}
// Default case: return full IP data
sendJSONResponse(w, data, http.StatusOK)
}
// Sends a JSON response with the given data and status code
func sendJSONResponse(w http.ResponseWriter, data any, statusCode int) {
w.Header().Set("Content-Type", "application/json; charset=utf-8")
w.WriteHeader(statusCode)
@@ -203,21 +228,16 @@ func sendJSONResponse(w http.ResponseWriter, data any, statusCode int) {
}
}
// Sends a JSON error response
func sendJSONError(w http.ResponseWriter, errMsg string, statusCode int) {
sendJSONResponse(w, map[string]string{"error": errMsg}, statusCode)
}
// Initializes and starts the server with the given GeoIPManager
func StartServer(ctx context.Context, geoIP *db.GeoIPManager) error {
// Create new server with the GeoIPManager
server := NewServer(geoIP)
// Set up signal handling for graceful shutdown
sigCh := make(chan os.Signal, 1)
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
// Handle shutdown signals
go func() {
select {
case <-sigCh:
@@ -229,6 +249,5 @@ func StartServer(ctx context.Context, geoIP *db.GeoIPManager) error {
}
}()
// Start the server
return server.Start(ctx)
}