cozystack/internal/observabilityproxy/server.go

/*
Copyright 2026 The Cozystack Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package observabilityproxy

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"math"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

// tailChunkBytes bounds the memory one tail holds: a longer line is passed
// on in pieces and flushed once its end arrives.
const tailChunkBytes = 64 << 10

// errTailIdle cancels a tail that received nothing for the idle timeout;
// it wraps DeadlineExceeded so a store that never answered reads as 504.
var errTailIdle = fmt.Errorf("tail idle: %w", context.DeadlineExceeded)

// errAborted ends a response whose status line already went out by
// cutting the connection, so a client sees a broken body instead of a
// complete one with lines missing.
var errAborted = &HTTPError{Code: http.StatusBadGateway, Message: "response aborted"}

// maxFormBytes bounds a POST body; LogsQL queries are a few KiB at most.
const maxFormBytes = 1 << 20

// Config is the runtime configuration of the proxy handler.
type Config struct {
	ClusterDomain string
	// UpstreamURL maps a store namespace to the vlselect base URL; nil
	// means the in-cluster Service of the "generic" logs storage.
	UpstreamURL     func(store string) string
	Logs            LogsLimits
	Limiter         LimiterConfig
	RequestTimeout  time.Duration
	TailIdleTimeout time.Duration
}

// StoreLookup resolves a request namespace to the namespace of its store.
type StoreLookup interface {
	Resolve(namespace string) (string, error)
}

// Server serves the tenant observability roots.
type Server struct {
	cfg      Config
	auth     *Authorizer
	stores   StoreLookup
	client   *http.Client
	limits   *limiter
	now      func() time.Time
	requests *prometheus.CounterVec
}

func NewServer(cfg Config, auth *Authorizer, stores StoreLookup, reg prometheus.Registerer) *Server {
	if cfg.UpstreamURL == nil {
		domain := cfg.ClusterDomain
		cfg.UpstreamURL = func(store string) string {
			return fmt.Sprintf("http://vlselect-generic.%s.svc.%s:9471", store, domain)
		}
	}
	s := &Server{
		cfg:    cfg,
		auth:   auth,
		stores: stores,
		// No client timeout: it would cut live tails. Non-tail requests are
		// bounded by a context deadline instead.
		client: &http.Client{Transport: http.DefaultTransport.(*http.Transport).Clone()},
		limits: newLimiter(cfg.Limiter),
		now:    time.Now,
		requests: prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: "observability_proxy_requests_total",
			Help: "Requests served, by route kind and status code.",
		}, []string{"route", "code"}),
	}
	if reg != nil {
		reg.MustRegister(s.requests, prometheus.NewGaugeFunc(prometheus.GaugeOpts{
			Name: "observability_proxy_open_tails",
			Help: "Live tails currently open.",
		}, func() float64 { return float64(s.limits.openTails()) }))
	}
	return s
}

type statusRecorder struct {
	http.ResponseWriter
	code int
}

func (r *statusRecorder) WriteHeader(code int) {
	if r.code == 0 {
		r.code = code
	}
	r.ResponseWriter.WriteHeader(code)
}

func (r *statusRecorder) Write(b []byte) (int, error) {
	if r.code == 0 {
		r.code = http.StatusOK
	}
	return r.ResponseWriter.Write(b)
}

func (r *statusRecorder) Unwrap() http.ResponseWriter { return r.ResponseWriter }

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	rec := &statusRecorder{ResponseWriter: w}
	route, herr := ParseRoute(r.URL.EscapedPath())
	if herr == nil {
		herr = s.serve(rec, r, route)
	}
	if herr != nil && herr != errAborted {
		writeError(rec, herr)
	}
	s.requests.WithLabelValues(route.Kind.String(), strconv.Itoa(rec.code)).Inc()
	if herr == errAborted {
		panic(http.ErrAbortHandler)
	}
}

func writeError(w http.ResponseWriter, e *HTTPError) {
	if e.Code == http.StatusUnauthorized {
		w.Header().Set("WWW-Authenticate", `Bearer realm="cozystack"`)
	}
	http.Error(w, e.Message, e.Code)
}

// serve answers an unimplemented root before authenticating, so later
// waves only add a case here.
func (s *Server) serve(w http.ResponseWriter, r *http.Request, route Route) *HTTPError {
	if route.Kind != NamespaceLogs {
		return errorf(http.StatusNotFound, "not found")
	}
	user, herr := s.authorize(r, route)
	if herr != nil {
		return herr
	}
	if r.Method != http.MethodGet && r.Method != http.MethodPost {
		w.Header().Set("Allow", "GET, POST")
		return errorf(http.StatusMethodNotAllowed, "method not allowed")
	}
	if wait, ok := s.limits.allowRate(user); !ok {
		w.Header().Set("Retry-After", strconv.FormatInt(int64(math.Ceil(wait.Seconds())), 10))
		return errorf(http.StatusTooManyRequests, "request rate limit exceeded")
	}
	endpoint, ok := strings.CutPrefix(route.Endpoint, "select/logsql/")
	if !ok {
		return errorf(http.StatusNotFound, "not found")
	}
	tail := endpoint == "tail"
	acquire := s.limits.acquire
	if tail {
		acquire = s.limits.acquireTail
	}
	release, ok := acquire(user)
	if !ok {
		return errorf(http.StatusTooManyRequests, "too many concurrent requests")
	}
	defer release()
	if herr := s.readForm(w, r); herr != nil {
		return herr
	}
	params, herr := BuildLogsParams(endpoint, r.Form, route.Namespace, s.now(), s.cfg.Logs)
	if herr != nil {
		return herr
	}

	store, err := s.stores.Resolve(route.Namespace)
	if err != nil {
		slog.Error("resolving log store", "namespace", route.Namespace, "error", err)
		return errorf(http.StatusServiceUnavailable, "cannot resolve the log store")
	}

	ctx, cancel := context.WithCancelCause(r.Context())
	defer cancel(nil)
	limit := s.cfg.RequestTimeout
	var idle *time.Timer
	if tail {
		// Started before the request, so a store that never sends the
		// headers is bounded as well as one that stops sending lines.
		if limit = s.cfg.TailIdleTimeout; limit > 0 {
			idle = time.AfterFunc(limit, func() { cancel(errTailIdle) })
			defer idle.Stop()
		}
	} else if limit > 0 {
		var stop context.CancelFunc
		ctx, stop = context.WithTimeout(ctx, limit)
		defer stop()
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		s.cfg.UpstreamURL(store)+"/select/logsql/"+endpoint, strings.NewReader(params.Encode()))
	if err != nil {
		return errorf(http.StatusInternalServerError, "building upstream request: %v", err)
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("AccountID", "0")
	req.Header.Set("ProjectID", "0")
	if tail {
		// vlselect holds even the headers of a gzip-encoded tail until the
		// first line; an explicit value stops the transport from asking for gzip.
		req.Header.Set("Accept-Encoding", "identity")
	}
	resp, err := s.client.Do(req)
	switch {
	case err == nil:
	case r.Context().Err() != nil:
		return errorf(http.StatusBadGateway, "client went away")
	case errors.Is(context.Cause(ctx), context.DeadlineExceeded):
		return errorf(http.StatusGatewayTimeout, "log store did not answer within %s", limit)
	default:
		slog.Warn("log store unreachable", "store", store, "error", err)
		return errorf(http.StatusBadGateway, "log store unreachable")
	}
	defer func() { _ = resp.Body.Close() }()

	if ct := resp.Header.Get("Content-Type"); ct != "" {
		w.Header().Set("Content-Type", ct)
	}
	// Cancelling the upstream request does not interrupt a write to a
	// client that stopped reading; only a write deadline does. A tail
	// renews it before every line, so only a stalled client is cut; the
	// server clears it once the response is finished.
	rc := http.NewResponseController(w)
	arm := func() {}
	if tail && s.cfg.TailIdleTimeout > 0 {
		arm = func() { _ = rc.SetWriteDeadline(time.Now().Add(s.cfg.TailIdleTimeout)) }
	} else if deadline, ok := ctx.Deadline(); ok {
		_ = rc.SetWriteDeadline(deadline)
	}
	arm()
	w.WriteHeader(resp.StatusCode)
	if rc.Flush() != nil {
		return nil
	}
	if tail {
		s.streamLines(w, rc, resp.Body, idle, arm)
	} else if _, err := io.Copy(w, resp.Body); err != nil {
		slog.Debug("copying upstream response", "error", err)
		return errAborted
	}
	return nil
}

// readForm reads the body under the request timeout, so a client that
// sends it slowly holds its slot only that long. Once the body is read the
// deadline is lifted: the server's background read would take a later
// expiry for a client that went away and cancel a tail. After a failed
// read it stays, so the server's drain of the unread body is bounded too.
func (s *Server) readForm(w http.ResponseWriter, r *http.Request) *HTTPError {
	r.Body = http.MaxBytesReader(w, r.Body, maxFormBytes)
	rc := http.NewResponseController(w)
	if s.cfg.RequestTimeout > 0 {
		_ = rc.SetReadDeadline(time.Now().Add(s.cfg.RequestTimeout))
	}
	err := r.ParseForm()
	switch {
	case err == nil:
		_ = rc.SetReadDeadline(time.Time{})
		return nil
	case errors.Is(err, os.ErrDeadlineExceeded):
		return errorf(http.StatusRequestTimeout, "request body not received within %s", s.cfg.RequestTimeout)
	default:
		return errorf(http.StatusBadRequest, "invalid form: %v", err)
	}
}

func (s *Server) authorize(r *http.Request, route Route) (string, *HTTPError) {
	token, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer ")
	if !ok {
		return "", errorf(http.StatusUnauthorized, "missing bearer token")
	}
	user, err := s.auth.Authenticate(r.Context(), strings.TrimSpace(token))
	if errors.Is(err, ErrUnauthenticated) {
		return "", errorf(http.StatusUnauthorized, "invalid bearer token")
	}
	if err != nil {
		slog.Error("token review failed", "error", err)
		return "", errorf(http.StatusServiceUnavailable, "authentication unavailable")
	}
	allowed, err := s.auth.Authorize(r.Context(), user, route.ResourceAttributes())
	if err != nil {
		slog.Error("subject access review failed", "error", err)
		return "", errorf(http.StatusServiceUnavailable, "authorization unavailable")
	}
	if !allowed {
		return "", errorf(http.StatusForbidden, "forbidden")
	}
	return user.Username, nil
}

// streamLines flushes after every line so a live tail reaches the browser
// as it happens; each line pushes back idle, which cancels the upstream
// stream once no line has arrived for the idle timeout.
func (s *Server) streamLines(w http.ResponseWriter, rc *http.ResponseController, body io.Reader, idle *time.Timer, arm func()) {
	rd := bufio.NewReaderSize(body, tailChunkBytes)
	for {
		chunk, err := rd.ReadSlice('\n')
		if len(chunk) > 0 {
			if idle != nil {
				idle.Reset(s.cfg.TailIdleTimeout)
			}
			arm()
			if _, werr := w.Write(chunk); werr != nil {
				return
			}
			if errors.Is(err, bufio.ErrBufferFull) {
				continue
			}
			if ferr := rc.Flush(); ferr != nil {
				return
			}
		}
		if err != nil && !errors.Is(err, bufio.ErrBufferFull) {
			return
		}
	}
}