gists/projects/go-tools/goipgrep/internal/ipinfo/ipinfo.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
}