Instrumentation: Proxy status code correction and various improvements (#47473)

For a proxied request, e.g. Grafana's datasource or plugin proxy: If the request is cancelled, e.g. from the browser, the HTTP status code is now 499 Client closed request instead of 502 Bad gateway. If the request times out, e.g. takes longer time than allowed, the HTTP status code is now 504 Gateway timeout instead of 502 Bad gateway. This also means that request metrics and logs will get their status codes adjusted according to above. Fixes #46337 Fixes #46338
2022-04-11 13:17:08 +02:00
parent c0104db72d
commit 4bc582570e
7 changed files with 387 additions and 80 deletions
@@ -0,0 +1,146 @@
+package proxyutil
+
+import (
+	"context"
+	"errors"
+	"log"
+	"net/http"
+	"net/http/httputil"
+	"strings"
+	"time"
+
+	glog "github.com/grafana/grafana/pkg/infra/log"
+)
+
+// StatusClientClosedRequest A non-standard status code introduced by nginx
+// for the case when a client closes the connection while nginx is processing
+// the request.
+// https://httpstatus.in/499/
+const StatusClientClosedRequest = 499
+
+// ReverseProxyOption reverse proxy option to configure a httputil.ReverseProxy.
+type ReverseProxyOption func(*httputil.ReverseProxy)
+
+// NewReverseProxy creates a new httputil.ReverseProxy with sane default configuration.
+func NewReverseProxy(logger glog.Logger, director func(*http.Request), opts ...ReverseProxyOption) *httputil.ReverseProxy {
+	if logger == nil {
+		panic("logger cannot be nil")
+	}
+
+	if director == nil {
+		panic("director cannot be nil")
+	}
+
+	p := &httputil.ReverseProxy{
+		FlushInterval: time.Millisecond * 200,
+		ErrorHandler:  errorHandler(logger),
+		ErrorLog:      log.New(&logWrapper{logger: logger}, "", 0),
+		Director:      director,
+	}
+
+	for _, opt := range opts {
+		opt(p)
+	}
+
+	origDirector := p.Director
+	p.Director = wrapDirector(origDirector)
+
+	if p.ModifyResponse == nil {
+		// nolint:bodyclose
+		p.ModifyResponse = modifyResponse(logger)
+	} else {
+		modResponse := p.ModifyResponse
+		p.ModifyResponse = func(resp *http.Response) error {
+			if err := modResponse(resp); err != nil {
+				return err
+			}
+
+			// nolint:bodyclose
+			return modifyResponse(logger)(resp)
+		}
+	}
+
+	return p
+}
+
+// wrapDirector wraps a director and adds additional functionality.
+func wrapDirector(d func(*http.Request)) func(req *http.Request) {
+	return func(req *http.Request) {
+		d(req)
+		PrepareProxyRequest(req)
+
+		// Clear Origin and Referer to avoid CORS issues
+		req.Header.Del("Origin")
+		req.Header.Del("Referer")
+	}
+}
+
+// modifyResponse enforces certain constraints on http.Response.
+func modifyResponse(logger glog.Logger) func(resp *http.Response) error {
+	return func(resp *http.Response) error {
+		resp.Header.Del("Set-Cookie")
+		SetProxyResponseHeaders(resp.Header)
+		return nil
+	}
+}
+
+type timeoutError interface {
+	error
+	Timeout() bool
+}
+
+// errorHandler handles any errors happening while proxying a request and enforces
+// certain HTTP status based on the kind of error.
+// If client cancel/close the request we return 499 StatusClientClosedRequest.
+// If timeout happens while communicating with upstream server we return http.StatusGatewayTimeout.
+// If any other error we return http.StatusBadGateway.
+func errorHandler(logger glog.Logger) func(http.ResponseWriter, *http.Request, error) {
+	return func(w http.ResponseWriter, r *http.Request, err error) {
+		if errors.Is(err, context.Canceled) {
+			logger.Debug("Proxy request cancelled by client")
+			w.WriteHeader(StatusClientClosedRequest)
+			return
+		}
+
+		// nolint:errorlint
+		if timeoutErr, ok := err.(timeoutError); ok && timeoutErr.Timeout() {
+			logger.Error("Proxy request timed out", "err", err)
+			w.WriteHeader(http.StatusGatewayTimeout)
+			return
+		}
+
+		logger.Error("Proxy request failed", "err", err)
+		w.WriteHeader(http.StatusBadGateway)
+	}
+}
+
+type logWrapper struct {
+	logger glog.Logger
+}
+
+// Write writes log messages as bytes from proxy.
+func (lw *logWrapper) Write(p []byte) (n int, err error) {
+	withoutNewline := strings.TrimSuffix(string(p), "\n")
+	lw.logger.Error("Proxy request error", "error", withoutNewline)
+	return len(p), nil
+}
+
+func WithTransport(transport http.RoundTripper) ReverseProxyOption {
+	if transport == nil {
+		panic("transport cannot be nil")
+	}
+
+	return ReverseProxyOption(func(rp *httputil.ReverseProxy) {
+		rp.Transport = transport
+	})
+}
+
+func WithModifyResponse(fn func(*http.Response) error) ReverseProxyOption {
+	if fn == nil {
+		panic("fn cannot be nil")
+	}
+
+	return ReverseProxyOption(func(rp *httputil.ReverseProxy) {
+		rp.ModifyResponse = fn
+	})
+}
@@ -0,0 +1,157 @@
+package proxyutil
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/stretchr/testify/require"
+)
+
+func TestReverseProxy(t *testing.T) {
+	t.Run("When proxying a request should enforce request and response constraints", func(t *testing.T) {
+		var actualReq *http.Request
+		upstream := newUpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+			actualReq = req
+			http.SetCookie(w, &http.Cookie{Name: "test"})
+			w.WriteHeader(http.StatusOK)
+		}))
+		t.Cleanup(upstream.Close)
+		rec := httptest.NewRecorder()
+		req := httptest.NewRequest(http.MethodGet, upstream.URL, nil)
+		req.Header.Set("X-Forwarded-Host", "forwarded.host.com")
+		req.Header.Set("X-Forwarded-Port", "8080")
+		req.Header.Set("X-Forwarded-Proto", "https")
+		req.Header.Set("Origin", "test.com")
+		req.Header.Set("Referer", "https://test.com/api")
+		req.RemoteAddr = "10.0.0.1"
+
+		rp := NewReverseProxy(log.New("test"), func(req *http.Request) {
+			req.Header.Set("X-KEY", "value")
+		})
+		require.NotNil(t, rp)
+		require.NotNil(t, rp.ModifyResponse)
+		rp.ServeHTTP(rec, req)
+
+		require.NotNil(t, actualReq)
+		require.Empty(t, actualReq.Header.Get("X-Forwarded-Host"))
+		require.Empty(t, actualReq.Header.Get("X-Forwarded-Port"))
+		require.Empty(t, actualReq.Header.Get("X-Forwarded-Proto"))
+		require.Equal(t, "10.0.0.1", actualReq.Header.Get("X-Forwarded-For"))
+		require.Empty(t, actualReq.Header.Get("Origin"))
+		require.Empty(t, actualReq.Header.Get("Referer"))
+		require.Equal(t, "value", actualReq.Header.Get("X-KEY"))
+		resp := rec.Result()
+		require.Empty(t, resp.Cookies())
+		require.Equal(t, "sandbox", resp.Header.Get("Content-Security-Policy"))
+		require.NoError(t, resp.Body.Close())
+	})
+
+	t.Run("When proxying a request using WithModifyResponse should call it before default ModifyResponse func", func(t *testing.T) {
+		var actualReq *http.Request
+		upstream := newUpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+			actualReq = req
+			http.SetCookie(w, &http.Cookie{Name: "test"})
+			w.WriteHeader(http.StatusOK)
+		}))
+		t.Cleanup(upstream.Close)
+		rec := httptest.NewRecorder()
+		req := httptest.NewRequest(http.MethodGet, upstream.URL, nil)
+		rp := NewReverseProxy(
+			log.New("test"),
+			func(req *http.Request) {
+				req.Header.Set("X-KEY", "value")
+			},
+			WithModifyResponse(func(r *http.Response) error {
+				r.Header.Set("X-KEY2", "value2")
+				return nil
+			}),
+		)
+		require.NotNil(t, rp)
+		require.NotNil(t, rp.ModifyResponse)
+		rp.ServeHTTP(rec, req)
+
+		require.NotNil(t, actualReq)
+		require.Equal(t, "value", actualReq.Header.Get("X-KEY"))
+		resp := rec.Result()
+		require.Empty(t, resp.Cookies())
+		require.Equal(t, "sandbox", resp.Header.Get("Content-Security-Policy"))
+		require.Equal(t, "value2", resp.Header.Get("X-KEY2"))
+		require.NoError(t, resp.Body.Close())
+	})
+
+	t.Run("Error handling should convert status codes depending on what kind of error it is", func(t *testing.T) {
+		timedOutTransport := http.DefaultTransport.(*http.Transport)
+		timedOutTransport.ResponseHeaderTimeout = time.Nanosecond
+
+		testCases := []struct {
+			desc               string
+			transport          http.RoundTripper
+			expectedStatusCode int
+		}{
+			{
+				desc:               "Cancelled request should return 499 Client closed request",
+				transport:          &cancelledRoundTripper{},
+				expectedStatusCode: StatusClientClosedRequest,
+			},
+			{
+				desc:               "Timed out request should return 504 Gateway timeout",
+				transport:          timedOutTransport,
+				expectedStatusCode: http.StatusGatewayTimeout,
+			},
+			{
+				desc:               "Failed request should return 502 Bad gateway",
+				transport:          &failingRoundTripper{},
+				expectedStatusCode: http.StatusBadGateway,
+			},
+		}
+
+		for _, tc := range testCases {
+			t.Run(tc.desc, func(t *testing.T) {
+				upstream := newUpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+					w.WriteHeader(http.StatusOK)
+				}))
+				t.Cleanup(upstream.Close)
+				rec := httptest.NewRecorder()
+				req := httptest.NewRequest(http.MethodGet, upstream.URL, nil)
+
+				rp := NewReverseProxy(
+					log.New("test"),
+					func(req *http.Request) {},
+					WithTransport(tc.transport),
+				)
+				require.NotNil(t, rp)
+				require.NotNil(t, rp.Transport)
+				require.Same(t, tc.transport, rp.Transport)
+				rp.ServeHTTP(rec, req)
+
+				resp := rec.Result()
+				require.Equal(t, tc.expectedStatusCode, resp.StatusCode)
+				require.NoError(t, resp.Body.Close())
+			})
+		}
+	})
+}
+
+func newUpstreamServer(t *testing.T, handler http.Handler) *httptest.Server {
+	t.Helper()
+
+	upstream := httptest.NewServer(handler)
+	return upstream
+}
+
+type cancelledRoundTripper struct{}
+
+func (cancelledRoundTripper) RoundTrip(*http.Request) (*http.Response, error) {
+	return nil, context.Canceled
+}
+
+type failingRoundTripper struct{}
+
+func (failingRoundTripper) RoundTrip(*http.Request) (*http.Response, error) {
+	return nil, errors.New("some error")
+}