132 lines
2.4 KiB
Go
132 lines
2.4 KiB
Go
package ipinfo
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/url"
|
|
"runtime"
|
|
"sort"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type Info struct {
|
|
IP string `json:"ip"`
|
|
Hostname string `json:"hostname,omitempty"`
|
|
City string `json:"city,omitempty"`
|
|
Region string `json:"region,omitempty"`
|
|
Country string `json:"country,omitempty"`
|
|
Org string `json:"org,omitempty"`
|
|
}
|
|
|
|
type Client struct {
|
|
Token string
|
|
Timeout time.Duration
|
|
}
|
|
|
|
type Cache interface {
|
|
Get(ip string) (Info, bool)
|
|
Put(ip string, info Info)
|
|
}
|
|
|
|
func (c Client) Lookup(ctx context.Context, ip string) (Info, error) {
|
|
u := fmt.Sprintf("https://ipinfo.io/%s/json", url.PathEscape(ip))
|
|
if c.Token != "" {
|
|
u += "?token=" + url.QueryEscape(c.Token)
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
|
|
if err != nil {
|
|
return Info{}, err
|
|
}
|
|
req.Header.Set("Accept", "application/json")
|
|
|
|
hc := &http.Client{Timeout: c.Timeout}
|
|
resp, err := hc.Do(req)
|
|
if err != nil {
|
|
return Info{}, err
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode != 200 {
|
|
return Info{}, fmt.Errorf("ipinfo: %s", resp.Status)
|
|
}
|
|
var info Info
|
|
if err := json.NewDecoder(resp.Body).Decode(&info); err != nil {
|
|
return Info{}, err
|
|
}
|
|
if info.IP == "" {
|
|
// Some error payloads decode but don't include the IP.
|
|
info.IP = ip
|
|
}
|
|
return info, nil
|
|
}
|
|
|
|
func LookupAll(ctx context.Context, ips []string, client Client, store Cache, jobs int) []Info {
|
|
if jobs <= 0 {
|
|
jobs = runtime.GOMAXPROCS(0)
|
|
}
|
|
in := make(chan string)
|
|
out := make(chan Info)
|
|
var wg sync.WaitGroup
|
|
|
|
worker := func() {
|
|
defer wg.Done()
|
|
for ip := range in {
|
|
if store != nil {
|
|
if info, ok := store.Get(ip); ok {
|
|
select {
|
|
case out <- info:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
continue
|
|
}
|
|
}
|
|
|
|
cctx, cancel := context.WithTimeout(ctx, client.Timeout)
|
|
info, err := client.Lookup(cctx, ip)
|
|
cancel()
|
|
if err == nil && info.IP != "" {
|
|
if store != nil {
|
|
store.Put(ip, info)
|
|
}
|
|
select {
|
|
case out <- info:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
wg.Add(jobs)
|
|
for i := 0; i < jobs; i++ {
|
|
go worker()
|
|
}
|
|
|
|
go func() {
|
|
defer close(in)
|
|
for _, ip := range ips {
|
|
select {
|
|
case in <- ip:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
go func() {
|
|
wg.Wait()
|
|
close(out)
|
|
}()
|
|
|
|
var res []Info
|
|
for info := range out {
|
|
res = append(res, info)
|
|
}
|
|
sort.Slice(res, func(i, j int) bool { return res[i].IP < res[j].IP })
|
|
return res
|
|
}
|