telemt-api/internal/aggregate/handlers.go

package aggregate

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/telemt/telemt-api/internal/config"
	"github.com/telemt/telemt-api/internal/geoip"
)

const pathPrefix = "/api/agg"

var aggUsernameRe = regexp.MustCompile(`^[A-Za-z0-9_.-]+$`)

type cacheEntry struct {
	body    []byte
	expires time.Time
}

// Handler serves GET /api/agg/* aggregate endpoints.
type Handler struct {
	Parsed   *config.Parsed
	Client   *http.Client
	Geo      *geoip.Service
	CacheTTL time.Duration

	cacheMu sync.Mutex
	cache   map[string]cacheEntry
}

// ResolveAliasesForLive resolves aliases for external endpoints (SSE/live).
func (h *Handler) ResolveAliasesForLive(r *http.Request) ([]string, error) {
	return h.resolveAliases(r)
}

// NewHandler builds an aggregate handler; client must use a non-nil Transport (e.g. gateway shared transport).
// Geo may be nil (no GeoLite2 lookups). cacheTTL 0 disables response caching.
func NewHandler(p *config.Parsed, client *http.Client, geo *geoip.Service, cacheTTL time.Duration) *Handler {
	return &Handler{
		Parsed:   p,
		Client:   client,
		Geo:      geo,
		CacheTTL: cacheTTL,
		cache:    make(map[string]cacheEntry),
	}
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		w.WriteHeader(http.StatusMethodNotAllowed)
		_ = json.NewEncoder(w).Encode(errEnvelope("method_not_allowed", "only GET is allowed"))
		return
	}
	sub := strings.TrimPrefix(r.URL.Path, pathPrefix)
	sub = strings.TrimPrefix(sub, "/")

	if h.CacheTTL > 0 {
		key := r.URL.Path + "\x00" + r.URL.RawQuery
		now := time.Now()
		h.cacheMu.Lock()
		ent, hit := h.cache[key]
		if hit && now.Before(ent.expires) {
			body := ent.body
			h.cacheMu.Unlock()
			w.Header().Set("Content-Type", "application/json; charset=utf-8")
			_, _ = w.Write(body)
			return
		}
		h.cacheMu.Unlock()

		rec := httptest.NewRecorder()
		h.dispatch(rec, r, sub)
		body := rec.Body.Bytes()
		if rec.Code == http.StatusOK {
			h.cacheMu.Lock()
			h.cache[key] = cacheEntry{body: append([]byte(nil), body...), expires: now.Add(h.CacheTTL)}
			h.cacheMu.Unlock()
		}
		copyRecorderToResponse(rec, w)
		return
	}

	h.dispatch(w, r, sub)
}

func copyRecorderToResponse(rec *httptest.ResponseRecorder, w http.ResponseWriter) {
	for k, vv := range rec.Header() {
		for _, v := range vv {
			w.Header().Add(k, v)
		}
	}
	w.WriteHeader(rec.Code)
	_, _ = w.Write(rec.Body.Bytes())
}

func (h *Handler) dispatch(w http.ResponseWriter, r *http.Request, sub string) {
	switch {
	case sub == "summary":
		h.handleSummary(w, r)
	case sub == "traffic":
		h.handleTraffic(w, r)
	case sub == "unique-ips":
		h.handleUniqueIPs(w, r)
	case sub == "users":
		h.handleUsers(w, r)
	case sub == "fleet-status":
		h.handleFleetStatus(w, r)
	case sub == "incidents":
		h.handleIncidents(w, r)
	case strings.HasPrefix(sub, "user/"):
		username := strings.TrimPrefix(sub, "user/")
		if username == "" {
			writeNotFound(w, "missing username")
			return
		}
		h.handleUserOne(w, r, username)
	default:
		writeNotFound(w, "unknown aggregate path")
	}
}

func errEnvelope(code, msg string) map[string]any {
	return map[string]any{
		"ok": false,
		"error": map[string]string{
			"code":    code,
			"message": msg,
		},
	}
}

func (h *Handler) resolveAliases(r *http.Request) ([]string, error) {
	q := r.URL.Query().Get("aliases")
	if strings.TrimSpace(q) != "" {
		var out []string
		for _, p := range strings.Split(q, ",") {
			a := strings.TrimSpace(p)
			if a == "" {
				continue
			}
			if h.Parsed.ByAlias[a] == nil {
				return nil, &resolveError{msg: "unknown alias: " + a}
			}
			out = append(out, a)
		}
		if len(out) == 0 {
			return nil, &resolveError{msg: "aliases query produced empty list"}
		}
		return out, nil
	}
	cfg := h.Parsed.Config.Aggregate
	if cfg != nil && len(cfg.IncludeAliases) > 0 {
		for _, a := range cfg.IncludeAliases {
			a = strings.TrimSpace(a)
			if a == "" {
				continue
			}
			if h.Parsed.ByAlias[a] == nil {
				return nil, &resolveError{msg: "aggregate.include_aliases: unknown alias: " + a}
			}
		}
		out := make([]string, 0, len(cfg.IncludeAliases))
		for _, a := range cfg.IncludeAliases {
			a = strings.TrimSpace(a)
			if a == "" {
				continue
			}
			out = append(out, a)
		}
		return out, nil
	}
	out := make([]string, 0, len(h.Parsed.Config.Servers))
	for i := range h.Parsed.Config.Servers {
		out = append(out, h.Parsed.Config.Servers[i].Alias)
	}
	return out, nil
}

type resolveError struct {
	msg string
}

func (e *resolveError) Error() string { return e.msg }

func anyUpstreamFailed(results []ServerFetchResult) bool {
	for _, fr := range results {
		if !fr.OK {
			return true
		}
	}
	return false
}

func (h *Handler) handleSummary(w http.ResponseWriter, r *http.Request) {
	aliases, err := h.resolveAliases(r)
	if err != nil {
		writeBadRequest(w, err)
		return
	}
	topN := 10
	if v := r.URL.Query().Get("top_n"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n > 0 {
			topN = n
		}
	}
	ctx, cancel := context.WithTimeout(r.Context(), 60*time.Second)
	defer cancel()
	results := FetchStatsUsers(ctx, h.Client, h.Parsed, aliases)
	data := BuildSummary(results, topN)
	partial := anyUpstreamFailed(results)
	writeAggOK(w, partial, data)
}

func (h *Handler) handleTraffic(w http.ResponseWriter, r *http.Request) {
	aliases, err := h.resolveAliases(r)
	if err != nil {
		writeBadRequest(w, err)
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), 60*time.Second)
	defer cancel()
	results := FetchStatsUsers(ctx, h.Client, h.Parsed, aliases)
	data := BuildTraffic(results)
	partial := anyUpstreamFailed(results)
	writeAggOK(w, partial, data)
}

func (h *Handler) handleUniqueIPs(w http.ResponseWriter, r *http.Request) {
	aliases, err := h.resolveAliases(r)
	if err != nil {
		writeBadRequest(w, err)
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), 60*time.Second)
	defer cancel()
	results := FetchStatsUsers(ctx, h.Client, h.Parsed, aliases)
	data := BuildUniqueIPs(results)
	if h.Geo != nil && !strings.EqualFold(r.URL.Query().Get("geo"), "false") {
		EnrichUniqueIPsGeo(data, h.Geo)
	}
	partial := anyUpstreamFailed(results)
	writeAggOK(w, partial, data)
}

func (h *Handler) handleUsers(w http.ResponseWriter, r *http.Request) {
	aliases, err := h.resolveAliases(r)
	if err != nil {
		writeBadRequest(w, err)
		return
	}
	includeLinks := strings.EqualFold(r.URL.Query().Get("include_links"), "true")
	var minOct uint64
	if v := r.URL.Query().Get("min_total_megabytes"); v != "" {
		if n, err := strconv.ParseFloat(v, 64); err == nil && n > 0 {
			minOct = uint64(n * float64(mebibyte))
		}
	} else if v := r.URL.Query().Get("min_total_octets"); v != "" {
		if n, err := strconv.ParseUint(v, 10, 64); err == nil {
			minOct = n
		}
	}
	ctx, cancel := context.WithTimeout(r.Context(), 60*time.Second)
	defer cancel()
	results := FetchStatsUsers(ctx, h.Client, h.Parsed, aliases)
	data := BuildUsers(results, includeLinks, minOct)
	partial := anyUpstreamFailed(results)
	writeAggOK(w, partial, data)
}

func (h *Handler) handleFleetStatus(w http.ResponseWriter, r *http.Request) {
	aliases, err := h.resolveAliases(r)
	if err != nil {
		writeBadRequest(w, err)
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), 60*time.Second)
	defer cancel()
	data := FetchFleetStatus(ctx, h.Client, h.Parsed, aliases)
	partial := data.ServersFailed > 0
	writeAggOK(w, partial, data)
}

func (h *Handler) handleUserOne(w http.ResponseWriter, r *http.Request, username string) {
	if !aggUsernameRe.MatchString(username) {
		writeBadRequestString(w, "invalid username")
		return
	}
	aliases, err := h.resolveAliases(r)
	if err != nil {
		writeBadRequest(w, err)
		return
	}
	includeLinks := strings.EqualFold(r.URL.Query().Get("include_links"), "true")
	ctx, cancel := context.WithTimeout(r.Context(), 60*time.Second)
	defer cancel()
	results := FetchStatsUsers(ctx, h.Client, h.Parsed, aliases)
	row := BuildSingleUser(results, username, includeLinks)
	if row == nil {
		writeNotFound(w, "user not found on any upstream")
		return
	}
	partial := anyUpstreamFailed(results)
	writeAggOK(w, partial, row)
}

func (h *Handler) handleIncidents(w http.ResponseWriter, r *http.Request) {
	aliases, err := h.resolveAliases(r)
	if err != nil {
		writeBadRequest(w, err)
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), 60*time.Second)
	defer cancel()
	data, partial := BuildIncidents(ctx, h, aliases)
	writeAggOK(w, partial, data)
}

func writeBadRequest(w http.ResponseWriter, err error) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(http.StatusBadRequest)
	_ = json.NewEncoder(w).Encode(errEnvelope("bad_request", err.Error()))
}

func writeBadRequestString(w http.ResponseWriter, msg string) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(http.StatusBadRequest)
	_ = json.NewEncoder(w).Encode(errEnvelope("bad_request", msg))
}

func writeNotFound(w http.ResponseWriter, msg string) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(http.StatusNotFound)
	_ = json.NewEncoder(w).Encode(errEnvelope("not_found", msg))
}

func writeAggOK(w http.ResponseWriter, partial bool, data any) {
	env := map[string]any{
		"ok":           true,
		"data":         data,
		"generated_at": time.Now().UTC().Format(time.RFC3339Nano),
	}
	if partial {
		env["partial"] = true
	}
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	_ = json.NewEncoder(w).Encode(env)
}