Bring back the Pyroscope API (#591)

* Bring back the Pyroscope API

* lint

Author: cyriltovena

pkg/phlare/modules.go

@@ -205,6 +205,9 @@ func (f *Phlare) initQuerier() (services.Service, error) {
 		return nil, err
 	}
 
+	f.Server.HTTP.Handle("/pyroscope/render", util.AuthenticateUser(f.Cfg.MultitenancyEnabled).Wrap(http.HandlerFunc(querierSvc.RenderHandler)))
+	f.Server.HTTP.Handle("/pyroscope/label-values", util.AuthenticateUser(f.Cfg.MultitenancyEnabled).Wrap(http.HandlerFunc(querierSvc.LabelValuesHandler)))
+
 	sm, err := services.NewManager(querierSvc, worker)
 	if err != nil {
 		return nil, err

pkg/util/http.go

@@ -25,6 +25,8 @@ import (
 	"github.com/weaveworks/common/user"
 	"golang.org/x/net/http2"
 	"gopkg.in/yaml.v3"
+
+	"github.com/grafana/phlare/pkg/tenant"
 )
 
 var defaultTransport http.RoundTripper = &http2.Transport{
@@ -348,3 +350,22 @@ func WriteJSONResponse(w http.ResponseWriter, v interface{}) {
 	// Also this isn't internal error, but error communicating with client.
 	_, _ = w.Write(data)
 }
+
+// AuthenticateUser propagates the user ID from HTTP headers back to the request's context.
+// If on is false, it will inject the default tenant ID.
+func AuthenticateUser(on bool) middleware.Interface {
+	return middleware.Func(func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if !on {
+				next.ServeHTTP(w, r.WithContext(user.InjectOrgID(r.Context(), tenant.DefaultTenantID)))
+				return
+			}
+			_, ctx, err := user.ExtractOrgIDFromHTTPRequest(r)
+			if err != nil {
+				http.Error(w, err.Error(), http.StatusUnauthorized)
+				return
+			}
+			next.ServeHTTP(w, r.WithContext(ctx))
+		})
+	})
+}

pkg/util/http_test.go

@@ -1,11 +1,14 @@
 package util_test
 
 import (
+	"net/http"
 	"net/http/httptest"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
+	"github.com/grafana/phlare/pkg/tenant"
 	"github.com/grafana/phlare/pkg/util"
 )
 
@@ -18,3 +21,32 @@ func TestWriteTextResponse(t *testing.T) {
 	assert.Equal(t, "hello world", w.Body.String())
 	assert.Equal(t, "text/plain", w.Header().Get("Content-Type"))
 }
+
+func TestMultitenantMiddleware(t *testing.T) {
+	w := httptest.NewRecorder()
+	r := httptest.NewRequest("GET", "http://localhost:8080", nil)
+
+	// No org ID header.
+	m := util.AuthenticateUser(true).Wrap(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		id, err := tenant.ExtractTenantIDFromContext(r.Context())
+		require.NoError(t, err)
+		assert.Equal(t, "1", id)
+	}))
+	m.ServeHTTP(w, r)
+	assert.Equal(t, http.StatusUnauthorized, w.Code)
+
+	w = httptest.NewRecorder()
+	r.Header.Set("X-Scope-OrgID", "1")
+	m.ServeHTTP(w, r)
+	assert.Equal(t, http.StatusOK, w.Code)
+
+	// No org ID header without auth.
+	r = httptest.NewRequest("GET", "http://localhost:8080", nil)
+	m = util.AuthenticateUser(false).Wrap(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		id, err := tenant.ExtractTenantIDFromContext(r.Context())
+		require.NoError(t, err)
+		assert.Equal(t, tenant.DefaultTenantID, id)
+	}))
+	m.ServeHTTP(w, r)
+	assert.Equal(t, http.StatusOK, w.Code)
+}