mirror of
https://github.com/skidoodle/ipinfo.git
synced 2026-04-28 01:27:34 +02:00
asn lookup support
This commit is contained in:
+107
-13
@@ -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
@@ -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
@@ -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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user