mirror of
https://github.com/TecharoHQ/anubis.git
synced 2026-04-26 10:02:42 +00:00
feat(thoth): cached ip to asn checker
Signed-off-by: Xe Iaso <me@xeiaso.net>
This commit is contained in:
@@ -0,0 +1,58 @@
|
||||
package thoth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net/netip"
|
||||
|
||||
iptoasnv1 "github.com/TecharoHQ/thoth-proto/gen/techaro/thoth/iptoasn/v1"
|
||||
"github.com/gaissmai/bart"
|
||||
"google.golang.org/grpc"
|
||||
)
|
||||
|
||||
type IPToASNWithCache struct {
|
||||
next iptoasnv1.IpToASNServiceClient
|
||||
table *bart.Table[*iptoasnv1.LookupResponse]
|
||||
}
|
||||
|
||||
func NewIpToASNWithCache(next iptoasnv1.IpToASNServiceClient) *IPToASNWithCache {
|
||||
return &IPToASNWithCache{
|
||||
next: next,
|
||||
table: &bart.Table[*iptoasnv1.LookupResponse]{},
|
||||
}
|
||||
}
|
||||
|
||||
func (ip2asn *IPToASNWithCache) Lookup(ctx context.Context, lr *iptoasnv1.LookupRequest, opts ...grpc.CallOption) (*iptoasnv1.LookupResponse, error) {
|
||||
addr, err := netip.ParseAddr(lr.GetIpAddress())
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("input is not an IP address: %w", err)
|
||||
}
|
||||
|
||||
cachedResponse, ok := ip2asn.table.Lookup(addr)
|
||||
if ok {
|
||||
return cachedResponse, nil
|
||||
}
|
||||
|
||||
resp, err := ip2asn.next.Lookup(ctx, lr, opts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var errs []error
|
||||
for _, cidr := range resp.GetCidr() {
|
||||
pfx, err := netip.ParsePrefix(cidr)
|
||||
if err != nil {
|
||||
errs = append(errs, err)
|
||||
continue
|
||||
}
|
||||
ip2asn.table.Insert(pfx, resp)
|
||||
}
|
||||
|
||||
if len(errs) != 0 {
|
||||
slog.Error("errors parsing IP prefixes", "err", errors.Join(errs...))
|
||||
}
|
||||
|
||||
return resp, nil
|
||||
}
|
||||
Reference in New Issue
Block a user