upstream.go
Rama master · réplica de lectura del repositorio original.
package aastro
import (
"bytes"
"context"
"errors"
"fmt"
"io"
"math"
"net"
"net/http"
"regexp"
"slices"
"strconv"
"strings"
"sync/atomic"
"time"
"github.com/go-chi/chi/v5"
"go.opentelemetry.io/otel"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/codes"
"go.opentelemetry.io/otel/propagation"
"go.opentelemetry.io/otel/trace"
"go.uber.org/zap"
"github.com/voidrunner3074/aastro/internal/circuitbreaker"
"github.com/voidrunner3074/aastro/internal/metric"
)
const methodQuery = "QUERY"
type upstream interface {
name() string
call(ctx context.Context, original *http.Request, originalBody []byte) *upstreamResponse
}
type upstreamResponse struct {
status int
headers http.Header
body []byte
err *upstreamError
}
type httpUpstream struct {
cfg upstreamConfig
state upstreamState
circuitBreaker *circuitbreaker.CircuitBreaker
metrics *metric.Metrics
log *zap.Logger
client *http.Client
streamClient *http.Client
}
type upstreamConfig struct {
id string
name string
hosts []string
path string
method string
timeout time.Duration
forwardHeaders []string
forwardQueries []string
forwardParams []string
trustedProxies []*net.IPNet
lbMode lbMode
policy upstreamPolicy
}
type upstreamState struct {
currentHostIdx uint64
activeConnections []int64
}
func (u *httpUpstream) name() string { return u.cfg.name }
// call executes the request with retries and circuit breaker protection,
// then validates the final response against the upstream's own policy.
func (u *httpUpstream) call(ctx context.Context, original *http.Request, originalBody []byte) *upstreamResponse {
log := u.log.With(
zap.String("upstream", u.cfg.name),
zap.String("request_id", requestIDFromContext(original.Context())),
)
resp := u.callWithRetry(ctx, original, originalBody, log)
u.updateCircuitBreaker(resp, log)
// Policy is applied after the circuit breaker update intentionally:
// a misconfigured requireBody should not cause the breaker to open.
u.applyPolicy(ctx, resp)
return resp
}
func (u *httpUpstream) callWithRetry(ctx context.Context, original *http.Request, originalBody []byte, log *zap.Logger) *upstreamResponse {
retry := u.cfg.policy.retry
// Same fallback as newRequest/proxy: an unset upstream method means
// "use the flow's own method". shouldRetry must judge idempotency
// against that effective method, not the raw (possibly empty) config
// value - otherwise retries never fire for the common case of an
// upstream that doesn't override the method.
method := u.cfg.method
if method == "" {
method = original.Method
}
var resp *upstreamResponse
for attempt := 0; attempt <= retry.maxRetries; attempt++ {
if attempt > 0 {
log.Debug("retrying upstream call",
zap.Int("attempt", attempt),
zap.Int("max_retries", retry.maxRetries),
zap.Int("previous_status", resp.status),
)
u.metrics.IncUpstreamRetriesTotal(routeFromContext(ctx), u.cfg.name)
}
if err := ctx.Err(); err != nil {
return &upstreamResponse{err: &upstreamError{kind: upstreamCanceled, err: err}}
}
if u.circuitBreaker != nil && !u.circuitBreaker.Allow() {
log.Error("circuit breaker open, request denied")
u.metrics.SetCircuitBreakerState(u.cfg.name, float64(u.circuitBreaker.State()))
return &upstreamResponse{
err: &upstreamError{
kind: upstreamCircuitOpen,
err: errors.New("upstream circuit breaker is open"),
},
}
}
resp = u.doCall(ctx, original, originalBody, log)
if attempt == retry.maxRetries || !u.shouldRetry(method, resp, retry) {
break
}
if retry.backoffDelay > 0 {
select {
case <-time.After(retry.backoffDelay):
case <-ctx.Done():
return &upstreamResponse{err: &upstreamError{kind: upstreamCanceled, err: ctx.Err()}}
}
}
}
return resp
}
func (u *httpUpstream) shouldRetry(method string, resp *upstreamResponse, retry retryPolicy) bool {
if !isIdempotent(method) {
return false
}
if resp.err == nil {
return slices.Contains(retry.retryOnStatuses, resp.status)
}
switch resp.err.kind.props().retry {
case retryAlways:
return true
case retryByStatus:
return slices.Contains(retry.retryOnStatuses, resp.status)
case retryNever:
return false
default:
return false
}
}
func isIdempotent(method string) bool {
switch method {
case http.MethodGet, http.MethodHead, http.MethodOptions, http.MethodTrace, http.MethodPut, http.MethodDelete, methodQuery:
return true
default:
return false
}
}
func (u *httpUpstream) updateCircuitBreaker(resp *upstreamResponse, log *zap.Logger) {
if u.circuitBreaker == nil {
return
}
switch breakerOutcomeFor(resp.err) {
case breakerFailure:
log.Error("upstream request failed, recording circuit breaker failure")
u.circuitBreaker.OnFailure()
case breakerSuccess:
u.circuitBreaker.OnSuccess()
case breakerNoSignal:
// The request never reached the upstream - denied by the breaker itself,
// canceled by the client, or never built. Recording success here would
// let an open breaker reset itself with the very request it rejected.
}
u.metrics.SetCircuitBreakerState(u.cfg.name, float64(u.circuitBreaker.State()))
}
// applyPolicy validates the response against the upstream's own policy rules.
// Violations are merged into resp.err so the original error kind is preserved.
func (u *httpUpstream) applyPolicy(ctx context.Context, resp *upstreamResponse) {
if resp.err != nil && !isStatusFailure(resp.err.kind) {
return
}
var errs []error
if u.cfg.policy.requireBody && len(resp.body) == 0 {
errs = append(errs, errors.New("empty body not allowed by upstream policy"))
}
if len(errs) == 0 {
return
}
u.metrics.IncUpstreamErrorsTotal(routeFromContext(ctx), u.cfg.name, "policy_violation")
combined := errors.Join(errs...)
switch {
case resp.err == nil:
resp.err = &upstreamError{kind: upstreamPolicyViolation, err: combined}
case resp.err.kind == upstreamClientError, resp.err.kind == upstreamRedirect:
// A policy violation on top of an otherwise-propagatable client error/redirect
// means the response isn't worth forwarding - downgrade to a policy violation (502).
resp.err = &upstreamError{
kind: upstreamPolicyViolation,
err: errors.Join(resp.err.err, combined),
}
default:
// 5xx keeps bad_status: it must keep feeding the circuit breaker
resp.err.err = errors.Join(resp.err.err, combined)
}
}
func (u *httpUpstream) doCall(ctx context.Context, original *http.Request, originalBody []byte, log *zap.Logger) *upstreamResponse {
ctx, cancel := context.WithTimeout(ctx, u.cfg.timeout)
defer cancel()
selectedHost := u.selectHost(log)
if u.cfg.lbMode == lbModeLeastConns {
atomic.AddInt64(&u.state.activeConnections[selectedHost], 1)
defer atomic.AddInt64(&u.state.activeConnections[selectedHost], -1)
}
req, err := u.newRequest(ctx, original, originalBody, u.cfg.hosts[selectedHost], log)
if err != nil {
return &upstreamResponse{err: &upstreamError{kind: upstreamInternal, err: err}}
}
span := trace.SpanFromContext(ctx)
span.SetAttributes(
attribute.String("http.method", req.Method),
attribute.String("http.url", req.URL.String()),
attribute.String("server.address", req.URL.Host),
attribute.String("aastro.upstream.host", u.cfg.hosts[selectedHost]),
)
otel.GetTextMapPropagator().Inject(ctx, propagation.HeaderCarrier(req.Header))
httpResp, err := u.client.Do(req)
if err != nil {
log.Error("upstream request failed", zap.Error(err))
span.RecordError(err)
span.SetStatus(codes.Error, "upstream request failed")
return &upstreamResponse{err: &upstreamError{kind: u.classifyDoError(err), err: err}}
}
defer httpResp.Body.Close()
span.SetAttributes(attribute.Int("http.status_code", httpResp.StatusCode))
body, uerr := u.readBody(ctx, httpResp.Body, log)
if uerr != nil {
span.RecordError(uerr.Unwrap())
span.SetStatus(codes.Error, uerr.Error())
return &upstreamResponse{status: httpResp.StatusCode, err: uerr}
}
resp := &upstreamResponse{
status: httpResp.StatusCode,
headers: u.filterHeaders(httpResp.Header),
body: body,
}
if kind, failed := u.classifyStatus(httpResp.StatusCode); failed {
if kind == upstreamBadStatus {
log.Error("upstream returned server error", zap.Int("status_code", httpResp.StatusCode))
span.SetStatus(codes.Error, http.StatusText(httpResp.StatusCode))
} else {
log.Debug(
"upstream returned non-success status",
zap.Int("status_code", httpResp.StatusCode),
zap.String("kind", string(kind)),
)
}
resp.err = &upstreamError{
kind: kind,
err: fmt.Errorf("upstream returned %d", httpResp.StatusCode),
}
}
return resp
}
func (u *httpUpstream) readBody(ctx context.Context, body io.ReadCloser, log *zap.Logger) ([]byte, *upstreamError) {
reader := io.Reader(body)
if u.cfg.policy.maxResponseBodySize > 0 {
log.Debug("applying response body size limit policy", zap.Int64("limit", u.cfg.policy.maxResponseBodySize))
// Read one byte past the limit: LimitReader signals exhaustion with io.EOF,
// which ReadAll reports as success. The extra byte is the only way to tell
// "body ended" from "body was truncated" - see the length check below
reader = io.LimitReader(reader, u.cfg.policy.maxResponseBodySize+1)
}
data, err := io.ReadAll(reader)
if err != nil {
if ctx.Err() != nil {
return nil, &upstreamError{kind: upstreamCanceled, err: ctx.Err()}
}
return nil, &upstreamError{kind: upstreamReadError, err: err}
}
if u.cfg.policy.maxResponseBodySize > 0 && int64(len(data)) > u.cfg.policy.maxResponseBodySize {
return nil, &upstreamError{
kind: upstreamBodyTooLarge,
err: fmt.Errorf("response size exceeds limit of %d bytes", u.cfg.policy.maxResponseBodySize),
}
}
return data, nil
}
func (u *httpUpstream) filterHeaders(headers http.Header) http.Header {
if len(u.cfg.policy.headerBlacklist) == 0 {
return headers.Clone()
}
filtered := make(http.Header, len(headers))
for header, values := range headers {
if _, blocked := u.cfg.policy.headerBlacklist[header]; !blocked {
filtered[header] = values
}
}
return filtered
}
func (u *httpUpstream) classifyStatus(status int) (upstreamErrorKind, bool) {
switch {
case status >= http.StatusInternalServerError:
return upstreamBadStatus, true
case status >= http.StatusBadRequest:
return upstreamClientError, true
case status >= http.StatusMultipleChoices:
return upstreamRedirect, true
default:
return "", false
}
}
func (u *httpUpstream) classifyDoError(err error) upstreamErrorKind {
switch {
case errors.Is(err, context.DeadlineExceeded):
return upstreamTimeout
case errors.Is(err, context.Canceled):
return upstreamCanceled
default:
return upstreamConnection
}
}
func (u *httpUpstream) newRequest(ctx context.Context, original *http.Request, originalBody []byte, targetHost string, log *zap.Logger) (*http.Request, error) {
path := expandPathParams(u.cfg.path, original)
path = strings.TrimPrefix(path, "/")
var hostPath string
if strings.HasSuffix(targetHost, "/") {
hostPath = targetHost + path
} else {
hostPath = targetHost + "/" + path
}
method := u.cfg.method
if method == "" {
method = original.Method
}
if method != http.MethodPost && method != http.MethodPut && method != http.MethodPatch && method != methodQuery {
originalBody = nil
}
target, err := http.NewRequestWithContext(ctx, method, hostPath, bytes.NewReader(originalBody))
if err != nil {
return nil, err
}
u.resolveQueries(target, original)
u.resolveHeaders(target, original, log)
return target, nil
}
var pathParamRegexp = regexp.MustCompile(`\{([^}]+)\}`)
func expandPathParams(path string, req *http.Request) string {
return pathParamRegexp.ReplaceAllStringFunc(path, func(match string) string {
name := match[1 : len(match)-1]
if value := chi.URLParam(req, name); value != "" {
return value
}
return match
})
}
func (u *httpUpstream) selectHost(log *zap.Logger) int64 {
if len(u.cfg.hosts) == 1 {
return 0
}
var selected int64
switch u.cfg.lbMode {
case lbModeRoundRobin:
idx := atomic.AddUint64(&u.state.currentHostIdx, 1)
// safe: idx % len is bounded by len(hosts) which is well within int64 range
selected = int64(idx % uint64(len(u.cfg.hosts))) //nolint:gosec // bounded by slice length
case lbModeLeastConns:
var minConns int64 = math.MaxInt64
for i := range u.cfg.hosts {
if c := atomic.LoadInt64(&u.state.activeConnections[i]); c < minConns {
minConns = c
selected = int64(i)
}
}
}
log.Debug("host selected",
zap.String("host", u.cfg.hosts[selected]),
zap.String("upstream", u.cfg.name),
)
return selected
}
func (u *httpUpstream) resolveQueries(target, original *http.Request) {
targetQ := target.URL.Query()
originalQ := original.URL.Query()
for _, pattern := range u.cfg.forwardQueries {
if pattern == "*" {
target.URL.RawQuery = originalQ.Encode()
return
}
if v := originalQ.Get(pattern); v != "" {
targetQ.Add(pattern, v)
}
}
for _, param := range u.cfg.forwardParams {
if param == "*" {
if rctx := chi.RouteContext(original.Context()); rctx != nil {
for i, key := range rctx.URLParams.Keys {
targetQ.Set(key, rctx.URLParams.Values[i])
}
}
target.URL.RawQuery = targetQ.Encode()
return
}
if v := chi.URLParam(original, param); v != "" {
targetQ.Set(param, v)
}
}
target.URL.RawQuery = targetQ.Encode()
}
func (u *httpUpstream) resolveHeaders(target, original *http.Request, log *zap.Logger) {
for _, pattern := range u.cfg.forwardHeaders {
if pattern == "*" {
target.Header = original.Header.Clone()
break
}
if strings.HasSuffix(pattern, "*") {
u.forwardHeadersByPrefix(target, original.Header, strings.TrimSuffix(pattern, "*"))
continue
}
if v := original.Header.Get(pattern); v != "" {
target.Header.Add(pattern, v)
}
}
target.Header.Set("Content-Type", original.Header.Get("Content-Type"))
clientIP := clientIPFromContext(original.Context())
if clientIP == "" {
clientIP = extractClientIP(original, u.cfg.trustedProxies)
}
parsedClientIP := net.ParseIP(clientIP)
if parsedClientIP == nil {
log.Warn("cannot parse client IP, skipping forward headers", zap.String("client_ip", clientIP))
return
}
isTLS := original.TLS != nil
port := u.resolvePort(u.effectiveHost(original), isTLS)
proto := "http"
if isTLS {
proto = "https"
}
if !u.isTrustedProxy(parsedClientIP) {
u.setUntrustedForwardingHeaders(original, target, clientIP, proto, port)
} else {
u.appendTrustedForwardingHeaders(original, target, clientIP, proto, port)
}
}
func (u *httpUpstream) forwardHeadersByPrefix(target *http.Request, src http.Header, prefix string) {
for name, values := range src {
if strings.HasPrefix(name, prefix) {
for _, v := range values {
target.Header.Add(name, v)
}
}
}
}
func (u *httpUpstream) setUntrustedForwardingHeaders(original, target *http.Request, clientIP, proto, port string) {
host := u.effectiveHost(original)
target.Header.Set("X-Forwarded-For", clientIP)
target.Header.Set("X-Forwarded-Proto", proto)
target.Header.Set("X-Forwarded-Host", host)
target.Header.Set("X-Forwarded-Port", port)
target.Header.Set("Forwarded", fmt.Sprintf("for=%s; proto=%s; host=%s", clientIP, proto, host))
}
func (u *httpUpstream) appendTrustedForwardingHeaders(original, target *http.Request, clientIP, proto, port string) {
if xff := original.Header.Get("X-Forwarded-For"); xff != "" {
target.Header.Set("X-Forwarded-For", xff+", "+clientIP)
} else {
target.Header.Set("X-Forwarded-For", clientIP)
}
if p := original.Header.Get("X-Forwarded-Proto"); p == "http" || p == "https" {
proto = p
}
target.Header.Set("X-Forwarded-Proto", proto)
host := u.effectiveHost(original)
if h := original.Header.Get("X-Forwarded-Host"); h != "" {
host = h
}
target.Header.Set("X-Forwarded-Host", host)
if p := original.Header.Get("X-Forwarded-Port"); isValidPort(p) {
target.Header.Set("X-Forwarded-Port", p)
} else {
target.Header.Set("X-Forwarded-Port", port)
}
hop := fmt.Sprintf("for=%s; proto=%s; host=%s", clientIP, proto, host)
if fwd := original.Header.Get("Forwarded"); fwd != "" {
target.Header.Set("Forwarded", fwd+", "+hop)
} else {
target.Header.Set("Forwarded", hop)
}
}
func (u *httpUpstream) isTrustedProxy(ip net.IP) bool {
for _, cidr := range u.cfg.trustedProxies {
if cidr.Contains(ip) {
return true
}
}
return false
}
func (u *httpUpstream) resolvePort(host string, isTLS bool) string {
_, port, err := net.SplitHostPort(host)
if err != nil {
if isTLS {
return "443"
}
return "80"
}
return port
}
func (u *httpUpstream) effectiveHost(req *http.Request) string {
if req.Host != "" {
return req.Host
}
return req.URL.Host
}
func isValidPort(p string) bool {
n, err := strconv.Atoi(p)
return err == nil && n >= 1 && n <= 65535
}
// proxy implements proxyCapable for streaming flows.
func (u *httpUpstream) proxy(ctx context.Context, w http.ResponseWriter, original *http.Request) error {
selectedHost := u.selectHost(u.log)
host := u.cfg.hosts[selectedHost]
path := expandPathParams(u.cfg.path, original)
path = strings.TrimPrefix(path, "/")
hostPath := host
if !strings.HasSuffix(host, "/") {
hostPath += "/"
}
hostPath += path
method := u.cfg.method
if method == "" {
method = original.Method
}
req, err := http.NewRequestWithContext(ctx, method, hostPath, original.Body)
if err != nil {
return fmt.Errorf("build upstream request: %w", err)
}
req.ContentLength = original.ContentLength
req.TransferEncoding = original.TransferEncoding
u.resolveQueries(req, original)
u.resolveHeaders(req, original, u.log)
span := trace.SpanFromContext(ctx)
span.SetAttributes(
attribute.String("http.method", req.Method),
attribute.String("http.url", req.URL.String()),
attribute.String("server.address", req.URL.Host),
attribute.String("aastro.upstream.host", host),
)
otel.GetTextMapPropagator().Inject(ctx, propagation.HeaderCarrier(req.Header))
resp, err := u.streamClient.Do(req) // #nosec G704
if err != nil {
span.RecordError(err)
span.SetStatus(codes.Error, "streaming upstream call failed")
return fmt.Errorf("upstream call: %w", err)
}
defer resp.Body.Close()
span.SetAttributes(attribute.Int("http.status_code", resp.StatusCode))
if resp.StatusCode >= http.StatusInternalServerError {
span.SetStatus(codes.Error, http.StatusText(resp.StatusCode))
}
for k, vv := range resp.Header {
if _, skip := hopByHopHeaders[k]; skip {
continue
}
for _, v := range vv {
w.Header().Add(k, v)
}
}
w.Header().Del("Content-Length")
w.WriteHeader(resp.StatusCode)
if err = streamCopy(w, resp.Body); err != nil {
if errors.Is(err, context.Canceled) {
span.SetAttributes(attribute.String("aastro.streaming.end_reason", "client_disconnect"))
return nil
}
span.RecordError(err)
span.SetStatus(codes.Error, "stream body failed")
return fmt.Errorf("stream body: %w", err)
}
return nil
}