routewatch/pkg/asinfo/asinfo.go

package asinfo

import (
	"bytes"
	"cmp"
	"compress/gzip"
	_ "embed"
	"encoding/json"
	"io"
	"slices"
	"strings"
	"sync"
)

//go:embed asdata.json.gz
var asDataGZ []byte

// Info represents information about an Autonomous System
type Info struct {
	ASN         int    `json:"asn"`
	Handle      string `json:"handle"`
	Description string `json:"description"`
}

// Registry provides access to AS information
type Registry struct {
	mu    sync.RWMutex
	byASN map[int]*Info
	// searchEntries holds every entry in AS number order, for Search.
	searchEntries []searchEntry
}

// searchEntry is an entry with its handle and description in lower case. They
// are put in lower case once, when the data loads, so that Search only
// compares.
type searchEntry struct {
	info             *Info
	lowerHandle      string
	lowerDescription string
}

var (
	//nolint:gochecknoglobals // Singleton pattern for embedded data
	defaultRegistry *Registry
	//nolint:gochecknoglobals // Singleton pattern for embedded data
	once sync.Once
)

// initRegistry initializes the default registry with embedded data
func initRegistry() {
	defaultRegistry = &Registry{
		byASN: make(map[int]*Info),
	}

	// Decompress the gzipped data
	gzReader, err := gzip.NewReader(bytes.NewReader(asDataGZ))
	if err != nil {
		panic("failed to create gzip reader: " + err.Error())
	}
	defer func() {
		_ = gzReader.Close()
	}()

	jsonData, err := io.ReadAll(gzReader)
	if err != nil {
		panic("failed to decompress AS data: " + err.Error())
	}

	var asInfos []Info

	err = json.Unmarshal(jsonData, &asInfos)
	if err != nil {
		panic("failed to unmarshal embedded AS data: " + err.Error())
	}

	slices.SortFunc(asInfos, func(a, b Info) int {
		return cmp.Compare(a.ASN, b.ASN)
	})

	defaultRegistry.searchEntries = make([]searchEntry, 0, len(asInfos))

	for i := range asInfos {
		info := &asInfos[i]
		defaultRegistry.byASN[info.ASN] = info
		defaultRegistry.searchEntries = append(defaultRegistry.searchEntries,
			searchEntry{
				info:             info,
				lowerHandle:      strings.ToLower(info.Handle),
				lowerDescription: strings.ToLower(info.Description),
			})
	}
}

// Get returns AS information for the given ASN
func Get(asn int) (*Info, bool) {
	once.Do(initRegistry)

	defaultRegistry.mu.RLock()
	defer defaultRegistry.mu.RUnlock()

	info, ok := defaultRegistry.byASN[asn]
	if !ok {
		return nil, false
	}

	// Return a copy to prevent mutation
	infoCopy := *info

	return &infoCopy, true
}

// GetDescription returns just the description for an ASN
func GetDescription(asn int) string {
	info, ok := Get(asn)
	if !ok {
		return ""
	}

	return info.Description
}

// GetHandle returns just the handle for an ASN
func GetHandle(asn int) string {
	info, ok := Get(asn)
	if !ok {
		return ""
	}

	return info.Handle
}

// Total returns the total number of AS entries in the registry
func Total() int {
	once.Do(initRegistry)

	defaultRegistry.mu.RLock()
	defer defaultRegistry.mu.RUnlock()

	return len(defaultRegistry.byASN)
}

// All returns all AS information as a slice
// Note: This creates a copy of all data, use sparingly
func All() []Info {
	once.Do(initRegistry)

	defaultRegistry.mu.RLock()
	defer defaultRegistry.mu.RUnlock()

	result := make([]Info, 0, len(defaultRegistry.byASN))
	for _, info := range defaultRegistry.byASN {
		result = append(result, *info)
	}

	return result
}

// Search finds the AS entries whose handle or description contains query,
// ignoring case. It returns the first limit of them in AS number order, and the
// number of entries that match. An empty query matches nothing.
func Search(query string, limit int) ([]Info, int) {
	once.Do(initRegistry)

	if query == "" {
		return nil, 0
	}

	query = strings.ToLower(query)

	defaultRegistry.mu.RLock()
	defer defaultRegistry.mu.RUnlock()

	var (
		results []Info
		total   int
	)

	for _, entry := range defaultRegistry.searchEntries {
		if strings.Contains(entry.lowerHandle, query) ||
			strings.Contains(entry.lowerDescription, query) {
			if total < limit {
				results = append(results, *entry.info)
			}

			total++
		}
	}

	return results, total
}