terdut-server/internal/prom/client.go

package prom

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math"
	"net"
	"net/http"
	"net/netip"
	"net/url"
	"strconv"
	"strings"
	"time"
)

const (
	// maxResponseBytes bounds what one query may make us read.
	maxResponseBytes = 5 << 20

	// maxSeries bounds how many series one query may return before we give up on
	// it; a rule expression without aggregation can match thousands.
	maxSeries = 2000

	// TargetPoints is how many points a graph is asked for.
	TargetPoints = 300
)

// ErrNotAllowed means the target host is not on the allowlist.
var ErrNotAllowed = errors.New("host is not in the allowed Prometheus hosts")

// Series is one time series, ready to chart. Points are [unix seconds, value].
type Series struct {
	Labels map[string]string `json:"labels"`
	Points [][2]float64      `json:"points"`
}

// Client queries Prometheus servers that the Allowlist permits.
type Client struct {
	http  *http.Client
	allow *Allowlist
}

// NewClient returns a client that only ever connects to allowed addresses. The
// check is made on the address the dialer is about to use, after name
// resolution, and redirects and environment proxies are off: either would let a
// permitted host send the request somewhere that is not.
func NewClient(allow *Allowlist, timeout time.Duration) *Client {
	if timeout <= 0 {
		timeout = 10 * time.Second
	}
	dialer := &net.Dialer{Timeout: timeout}
	tr := &http.Transport{
		Proxy:               nil,
		DialContext:         guardedDial(allow, dialer),
		TLSHandshakeTimeout: timeout,
		MaxIdleConns:        4,
		IdleConnTimeout:     30 * time.Second,
	}
	return &Client{
		allow: allow,
		http: &http.Client{
			Timeout:   timeout,
			Transport: tr,
			CheckRedirect: func(*http.Request, []*http.Request) error {
				return http.ErrUseLastResponse
			},
		},
	}
}

// Enabled reports whether any host is allowed at all.
func (c *Client) Enabled() bool { return c != nil && !c.allow.Empty() }

// Allowed is the early, resolution-free answer: whether base could be queried.
func (c *Client) Allowed(base *url.URL) bool {
	return c != nil && c.allow.MatchesHost(base)
}

func guardedDial(allow *Allowlist, d *net.Dialer) func(ctx context.Context, network, addr string) (net.Conn, error) {
	return func(ctx context.Context, network, addr string) (net.Conn, error) {
		host, port, err := net.SplitHostPort(addr)
		if err != nil {
			return nil, err
		}
		var ips []netip.Addr
		if ip, err := netip.ParseAddr(host); err == nil {
			ips = []netip.Addr{ip}
		} else {
			ips, err = net.DefaultResolver.LookupNetIP(ctx, "ip", host)
			if err != nil {
				return nil, err
			}
		}
		var lastErr error = ErrNotAllowed
		for _, ip := range ips {
			if !allow.permits(host, port, ip) {
				continue
			}
			// Dial the checked address itself, so a second lookup cannot return
			// something different.
			conn, err := d.DialContext(ctx, network, net.JoinHostPort(ip.String(), port))
			if err == nil {
				return conn, nil
			}
			lastErr = err
		}
		return nil, lastErr
	}
}

// QueryRange runs expr over [start, end] at the given step.
func (c *Client) QueryRange(ctx context.Context, base *url.URL, expr string, start, end time.Time, step time.Duration) ([]Series, error) {
	form := url.Values{
		"query": {expr},
		"start": {strconv.FormatInt(start.Unix(), 10)},
		"end":   {strconv.FormatInt(end.Unix(), 10)},
		"step":  {strconv.FormatInt(int64(step/time.Second), 10)},
	}
	res, err := c.do(ctx, base, "/api/v1/query_range", form)
	if err != nil {
		return nil, err
	}
	if res.ResultType != "matrix" {
		return nil, fmt.Errorf("unexpected result type %q", res.ResultType)
	}
	var raw []struct {
		Metric map[string]string    `json:"metric"`
		Values [][2]json.RawMessage `json:"values"`
	}
	if err := json.Unmarshal(res.Result, &raw); err != nil {
		return nil, fmt.Errorf("decode matrix: %w", err)
	}
	out := make([]Series, 0, len(raw))
	for _, r := range raw {
		s := Series{Labels: r.Metric, Points: make([][2]float64, 0, len(r.Values))}
		for _, v := range r.Values {
			if p, ok := point(v); ok {
				s.Points = append(s.Points, p)
			}
		}
		out = append(out, s)
	}
	return out, nil
}

// Query runs expr at a single instant and returns one point per series.
func (c *Client) Query(ctx context.Context, base *url.URL, expr string, at time.Time) ([]Series, error) {
	form := url.Values{
		"query": {expr},
		"time":  {strconv.FormatInt(at.Unix(), 10)},
	}
	res, err := c.do(ctx, base, "/api/v1/query", form)
	if err != nil {
		return nil, err
	}
	if res.ResultType != "vector" {
		return nil, fmt.Errorf("unexpected result type %q", res.ResultType)
	}
	var raw []struct {
		Metric map[string]string  `json:"metric"`
		Value  [2]json.RawMessage `json:"value"`
	}
	if err := json.Unmarshal(res.Result, &raw); err != nil {
		return nil, fmt.Errorf("decode vector: %w", err)
	}
	out := make([]Series, 0, len(raw))
	for _, r := range raw {
		s := Series{Labels: r.Metric}
		if p, ok := point(r.Value); ok {
			s.Points = [][2]float64{p}
		}
		out = append(out, s)
	}
	return out, nil
}

type apiResult struct {
	ResultType string          `json:"resultType"`
	Result     json.RawMessage `json:"result"`
}

func (c *Client) do(ctx context.Context, base *url.URL, apiPath string, form url.Values) (*apiResult, error) {
	if !c.Allowed(base) {
		return nil, ErrNotAllowed
	}
	u := *base
	u.Path = strings.TrimSuffix(base.Path, "/") + apiPath
	// POST, so a long expression is not limited by URL length.
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, u.String(), strings.NewReader(form.Encode()))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("Accept", "application/json")

	resp, err := c.http.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes+1))
	if err != nil {
		return nil, err
	}
	if len(body) > maxResponseBytes {
		return nil, errors.New("response too large")
	}
	var env struct {
		Status    string     `json:"status"`
		Data      *apiResult `json:"data"`
		ErrorType string     `json:"errorType"`
		Error     string     `json:"error"`
	}
	if err := json.Unmarshal(body, &env); err != nil {
		return nil, fmt.Errorf("prometheus answered %d with a body that is not JSON", resp.StatusCode)
	}
	if resp.StatusCode/100 != 2 || env.Status != "success" || env.Data == nil {
		msg := env.Error
		if msg == "" {
			msg = "status " + strconv.Itoa(resp.StatusCode)
		}
		return nil, fmt.Errorf("prometheus: %s", truncate(msg, 200))
	}
	if n := countResults(env.Data.Result); n > maxSeries {
		return nil, fmt.Errorf("query matches %d series, over the limit of %d", n, maxSeries)
	}
	return env.Data, nil
}

func countResults(raw json.RawMessage) int {
	var items []json.RawMessage
	if json.Unmarshal(raw, &items) != nil {
		return 0
	}
	return len(items)
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "…"
}

// point decodes a [timestamp, "value"] pair. Prometheus writes values as
// strings so that NaN and ±Inf survive; JSON cannot carry them, so those points
// are dropped.
func point(v [2]json.RawMessage) ([2]float64, bool) {
	var ts float64
	var s string
	if json.Unmarshal(v[0], &ts) != nil || json.Unmarshal(v[1], &s) != nil {
		return [2]float64{}, false
	}
	f, err := strconv.ParseFloat(s, 64)
	if err != nil || math.IsNaN(f) || math.IsInf(f, 0) {
		return [2]float64{}, false
	}
	return [2]float64{ts, f}, true
}

// StepFor picks a query step that yields about TargetPoints over the span, never
// finer than a second.
func StepFor(start, end time.Time) time.Duration {
	span := end.Sub(start)
	step := span / TargetPoints
	if step < time.Second {
		return time.Second
	}
	return step.Round(time.Second)
}