asn lookup support

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 {
+	var asnRecord db.ASNRecord
-		AutonomousSystemNumber       uint   `maxminddb:"autonomous_system_number"`
-		AutonomousSystemOrganization string `maxminddb:"autonomous_system_organization"`
-	}
 	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

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
+	cityDB       *maxminddb.Reader
-	asnDB  *maxminddb.Reader
+	asnDB        *maxminddb.Reader
-	mu     sync.RWMutex
+	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
 }

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.Handle("/health", utils.HealthCheck())
-	mux.HandleFunc("/", s.handler)             // Main handler
+	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) router(w http.ResponseWriter, r *http.Request) {
-func (s *Server) handler(w http.ResponseWriter, r *http.Request) {
-	requestedThings := strings.Split(strings.Trim(r.URL.Path, "/"), "/")
-
-	// Enable gzip compression if requested
 	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)
 }