From dfec9ba60213ad7e152060c48afce8416a9e310c Mon Sep 17 00:00:00 2001 From: skidoodle Date: Thu, 31 Jul 2025 19:18:32 +0200 Subject: [PATCH] asn lookup support --- internal/common/common.go | 120 +++++++++++++++++++++++++++++++++----- internal/db/db.go | 46 ++++++++++++++- internal/server/server.go | 81 +++++++++++++++---------- 3 files changed, 200 insertions(+), 47 deletions(-) diff --git a/internal/common/common.go b/internal/common/common.go index a30b55c..cbed767 100644 --- a/internal/common/common.go +++ b/internal/common/common.go @@ -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 diff --git a/internal/db/db.go b/internal/db/db.go index 9165b48..a40995c 100644 --- a/internal/db/db.go +++ b/internal/db/db.go @@ -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 } diff --git a/internal/server/server.go b/internal/server/server.go index b80c4ca..5924607 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -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/ or /AS.", 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) }