Fix multitenancy when using OTLP HTTP

Fixes #495 Signed-off-by: Goutham Veeramachaneni <[email protected]>
grafana · Oct 4, 2022 · 7839cd6 · 7839cd6
1 parent a1e620f
commit 7839cd6
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 3 deletions.
diff --git a/modules/distributor/receiver/middleware.go b/modules/distributor/receiver/middleware.go
@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/weaveworks/common/user"
+	"go.opentelemetry.io/collector/client"
 	"go.opentelemetry.io/collector/consumer"
 	"go.opentelemetry.io/collector/pdata/ptrace"
 
@@ -56,8 +57,15 @@ func (m *multiTenancyMiddleware) Wrap(next consumer.Traces) consumer.Traces {
 		var err error
 		_, ctx, err = user.ExtractFromGRPCRequest(ctx)
 		if err != nil {
-			log.Logger.Log("msg", "failed to extract org id", "err", err)
-			return err
+			// Maybe its a HTTP request.
+			info := client.FromContext(ctx)
+			orgIDs := info.Metadata.Get(user.OrgIDHeaderName)
+			if len(orgIDs) != 1 {
+				log.Logger.Log("msg", "failed to extract org id", "err", err)
+				return err
+			}
+
+			ctx = user.InjectOrgID(ctx, orgIDs[0])
 		}
 		return next.ConsumeTraces(ctx, td)
 	})

diff --git a/modules/distributor/receiver/middleware_test.go b/modules/distributor/receiver/middleware_test.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"github.com/weaveworks/common/user"
+	"go.opentelemetry.io/collector/client"
 	"go.opentelemetry.io/collector/consumer"
 	"go.opentelemetry.io/collector/pdata/ptrace"
 	"google.golang.org/grpc/metadata"
@@ -54,7 +55,7 @@ func TestFakeTenantMiddleware(t *testing.T) {
 func TestMultiTenancyMiddleware(t *testing.T) {
 	m := MultiTenancyMiddleware()
 
-	t.Run("injects org id", func(t *testing.T) {
+	t.Run("injects org id grpc", func(t *testing.T) {
 		tenantID := "test-tenant-id"
 
 		consumer := newAssertingConsumer(t, func(t *testing.T, ctx context.Context) {
@@ -70,6 +71,25 @@ func TestMultiTenancyMiddleware(t *testing.T) {
 		require.NoError(t, m.Wrap(consumer).ConsumeTraces(ctx, ptrace.Traces{}))
 	})
 
+	t.Run("injects org id http", func(t *testing.T) {
+		tenantID := "test-tenant-id"
+
+		consumer := newAssertingConsumer(t, func(t *testing.T, ctx context.Context) {
+			orgID, err := user.ExtractOrgID(ctx)
+			require.NoError(t, err)
+			require.Equal(t, orgID, tenantID)
+		})
+
+		info := client.Info{
+			Metadata: client.NewMetadata(map[string][]string{
+				"x-scope-OrgID": {tenantID},
+			}),
+		}
+
+		ctx := client.NewContext(context.Background(), info)
+		require.NoError(t, m.Wrap(consumer).ConsumeTraces(ctx, ptrace.Traces{}))
+	})
+
 	t.Run("returns error if org id cannot be extracted", func(t *testing.T) {
 		// no need to assert anything, because the wrapped function is never called
 		consumer := newAssertingConsumer(t, func(t *testing.T, ctx context.Context) {})

diff --git a/modules/distributor/receiver/shim.go b/modules/distributor/receiver/shim.go
@@ -184,6 +184,32 @@ func New(receiverCfg map[string]interface{}, pusher BatchPusher, middleware Midd
 			return nil, fmt.Errorf("receiver factory not found for type: %s", componentID.Type())
 		}
 
+		// Make sure that the headers are added to context. Required for Authentication.
+		switch componentID.Type() {
+		case "otlp":
+			otlpRecvCfg := cfg.(*otlpreceiver.Config)
+
+			if otlpRecvCfg.HTTP != nil {
+				otlpRecvCfg.HTTP.IncludeMetadata = true
+				cfg = otlpRecvCfg
+			}
+
+		case "zipkin":
+			zipkinRecvCfg := cfg.(*zipkinreceiver.Config)
+
+			zipkinRecvCfg.HTTPServerSettings.IncludeMetadata = true
+			cfg = zipkinRecvCfg
+
+		case "jaeger":
+			jaegerRecvCfg := cfg.(*jaegerreceiver.Config)
+
+			if jaegerRecvCfg.ThriftHTTP != nil {
+				jaegerRecvCfg.ThriftHTTP.IncludeMetadata = true
+			}
+
+			cfg = jaegerRecvCfg
+		}
+
 		receiver, err := factoryBase.CreateTracesReceiver(ctx, params, cfg, middleware.Wrap(shim))
 		if err != nil {
 			return nil, err