From 6c52c602bdc4ce62e1b12d5820d86e9c25ea054c Mon Sep 17 00:00:00 2001
From: Owen Diehl <ow.diehl@gmail.com>
Date: Thu, 12 Aug 2021 12:51:47 -0400
Subject: [PATCH] uses more fleshed out cortex auth utility & adds new
 auth-ignored routes

---
 pkg/loki/fake_auth.go | 42 ------------------------------------------
 pkg/loki/loki.go      | 26 ++++++++++++++------------
 2 files changed, 14 insertions(+), 54 deletions(-)
 delete mode 100644 pkg/loki/fake_auth.go

diff --git a/pkg/loki/fake_auth.go b/pkg/loki/fake_auth.go
deleted file mode 100644
index 05d1154cf45b..000000000000
--- a/pkg/loki/fake_auth.go
+++ /dev/null
@@ -1,42 +0,0 @@
-package loki
-
-import (
-	"context"
-	"net/http"
-
-	"github.com/weaveworks/common/middleware"
-	"github.com/weaveworks/common/user"
-	"google.golang.org/grpc"
-)
-
-// Fake auth middlewares just injects a fake userID, so the rest of the code
-// can continue to be multitenant.
-
-var fakeHTTPAuthMiddleware = middleware.Func(func(next http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		ctx := user.InjectOrgID(r.Context(), "fake")
-		next.ServeHTTP(w, r.WithContext(ctx))
-	})
-})
-
-var fakeGRPCAuthUnaryMiddleware = func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
-	ctx = user.InjectOrgID(ctx, "fake")
-	return handler(ctx, req)
-}
-
-var fakeGRPCAuthStreamMiddleware = func(srv interface{}, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
-	ctx := user.InjectOrgID(ss.Context(), "fake")
-	return handler(srv, serverStream{
-		ctx:          ctx,
-		ServerStream: ss,
-	})
-}
-
-type serverStream struct {
-	ctx context.Context
-	grpc.ServerStream
-}
-
-func (ss serverStream) Context() context.Context {
-	return ss.ctx
-}
diff --git a/pkg/loki/loki.go b/pkg/loki/loki.go
index f5aa891b0cd3..b68cc5bfa68a 100644
--- a/pkg/loki/loki.go
+++ b/pkg/loki/loki.go
@@ -15,6 +15,7 @@ import (
 	"github.com/grafana/loki/pkg/storage/stores/shipper/compactor"
 	"github.com/grafana/loki/pkg/validation"
 
+	"github.com/cortexproject/cortex/pkg/util/fakeauth"
 	"github.com/cortexproject/cortex/pkg/util/flagext"
 	"github.com/cortexproject/cortex/pkg/util/modules"
 	"github.com/prometheus/client_golang/prometheus"
@@ -46,7 +47,6 @@ import (
 	"github.com/grafana/loki/pkg/storage"
 	"github.com/grafana/loki/pkg/storage/chunk"
 	"github.com/grafana/loki/pkg/tracing"
-	serverutil "github.com/grafana/loki/pkg/util/server"
 )
 
 // Config is the root config for Loki.
@@ -223,17 +223,19 @@ func New(cfg Config) (*Loki, error) {
 }
 
 func (t *Loki) setupAuthMiddleware() {
-	t.Cfg.Server.GRPCMiddleware = []grpc.UnaryServerInterceptor{serverutil.RecoveryGRPCUnaryInterceptor}
-	t.Cfg.Server.GRPCStreamMiddleware = []grpc.StreamServerInterceptor{serverutil.RecoveryGRPCStreamInterceptor}
-	if t.Cfg.AuthEnabled {
-		t.Cfg.Server.GRPCMiddleware = append(t.Cfg.Server.GRPCMiddleware, middleware.ServerUserHeaderInterceptor)
-		t.Cfg.Server.GRPCStreamMiddleware = append(t.Cfg.Server.GRPCStreamMiddleware, GRPCStreamAuthInterceptor)
-		t.HTTPAuthMiddleware = middleware.AuthenticateUser
-	} else {
-		t.Cfg.Server.GRPCMiddleware = append(t.Cfg.Server.GRPCMiddleware, fakeGRPCAuthUnaryMiddleware)
-		t.Cfg.Server.GRPCStreamMiddleware = append(t.Cfg.Server.GRPCStreamMiddleware, fakeGRPCAuthStreamMiddleware)
-		t.HTTPAuthMiddleware = fakeHTTPAuthMiddleware
-	}
+	// Don't check auth header on TransferChunks, as we weren't originally
+	// sending it and this could cause transfers to fail on update.
+	t.HTTPAuthMiddleware = fakeauth.SetupAuthMiddleware(&t.Cfg.Server, t.Cfg.AuthEnabled,
+		// Also don't check auth for these gRPC methods, since single call is used for multiple users (or no user like health check).
+		[]string{
+			"/grpc.health.v1.Health/Check",
+			"/logproto.Ingester/TransferChunks",
+			"/frontend.Frontend/Process",
+			"/frontend.Frontend/NotifyClientShutdown",
+			"/schedulerpb.SchedulerForFrontend/FrontendLoop",
+			"/schedulerpb.SchedulerForQuerier/QuerierLoop",
+			"/schedulerpb.SchedulerForQuerier/NotifyQuerierShutdown",
+		})
 }
 
 var GRPCStreamAuthInterceptor = func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {