From eeb01126b35171451e0146f116d31d1797190c2f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Peter=20=C5=A0tibran=C3=BD?= <pstibrany@gmail.com>
Date: Wed, 25 Jun 2025 10:06:45 +0200
Subject: [PATCH] grpc: wrapContext should not replace existing metadata, but
 append to it. (#107126)

* grpc: wrapContext should not replace existing metadata, but append to it.

* Add test for wrapContext if there's no previous metadata.

* Rename encodeIdentityInMetadata to encodeIdentityInMetadataPairs
---
 .../unified/resource/grpc/authenticator.go    | 14 ++++----
 .../resource/grpc/authenticator_test.go       | 35 +++++++++++++++++--
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/pkg/storage/unified/resource/grpc/authenticator.go b/pkg/storage/unified/resource/grpc/authenticator.go
index 577e752ac2d..270835a85ef 100644
--- a/pkg/storage/unified/resource/grpc/authenticator.go
+++ b/pkg/storage/unified/resource/grpc/authenticator.go
@@ -11,8 +11,9 @@ import (
 	"google.golang.org/grpc/status"
 
 	"github.com/grafana/authlib/types"
-	"github.com/grafana/grafana/pkg/apimachinery/identity"
 	"go.opentelemetry.io/otel/trace"
+
+	"github.com/grafana/grafana/pkg/apimachinery/identity"
 )
 
 const (
@@ -139,15 +140,16 @@ func wrapContext(ctx context.Context) (context.Context, error) {
 	}
 
 	// set grpc metadata into the context to pass to the grpc server
-	return metadata.NewOutgoingContext(ctx, encodeIdentityInMetadata(user)), nil
+	ctx = metadata.AppendToOutgoingContext(ctx, encodeIdentityInMetadataPairs(user)...)
+	return ctx, nil
 }
 
-func encodeIdentityInMetadata(user identity.Requester) metadata.MD {
+func encodeIdentityInMetadataPairs(user identity.Requester) []string {
 	id, _ := user.GetInternalID()
 
-	logger.Debug("encodeIdentityInMetadata", "user.id", user.GetID(), "user.Login", user.GetLogin(), "user.Name", user.GetName())
+	logger.Debug("encodeIdentityInMetadataPairs", "user.id", user.GetID(), "user.Login", user.GetLogin(), "user.Name", user.GetName())
 
-	return metadata.Pairs(
+	return []string{
 		// This should be everything needed to recreate the user
 		mdToken, user.GetIDToken(),
 
@@ -161,5 +163,5 @@ func encodeIdentityInMetadata(user identity.Requester) metadata.MD {
 		// TODO, Remove after this is deployed to unified storage
 		"grafana-userid", strconv.FormatInt(id, 10),
 		"grafana-useruid", user.GetRawIdentifier(),
-	)
+	}
 }
diff --git a/pkg/storage/unified/resource/grpc/authenticator_test.go b/pkg/storage/unified/resource/grpc/authenticator_test.go
index 7a206b32820..69a4501edfe 100644
--- a/pkg/storage/unified/resource/grpc/authenticator_test.go
+++ b/pkg/storage/unified/resource/grpc/authenticator_test.go
@@ -1,13 +1,16 @@
 package grpc
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/metadata"
 
 	claims "github.com/grafana/authlib/types"
-	"github.com/grafana/grafana/pkg/apimachinery/identity"
 	"go.opentelemetry.io/otel/trace/noop"
+
+	"github.com/grafana/grafana/pkg/apimachinery/identity"
 )
 
 func TestBasicEncodeDecode(t *testing.T) {
@@ -22,8 +25,8 @@ func TestBasicEncodeDecode(t *testing.T) {
 
 	auth := &Authenticator{Tracer: noop.NewTracerProvider().Tracer("")}
 
-	md := encodeIdentityInMetadata(before)
-	after, err := auth.decodeMetadata(md)
+	md := encodeIdentityInMetadataPairs(before)
+	after, err := auth.decodeMetadata(metadata.Pairs(md...))
 	require.NoError(t, err)
 	require.Equal(t, before.GetID(), after.GetID())
 	require.Equal(t, before.GetUID(), after.GetUID())
@@ -33,3 +36,29 @@ func TestBasicEncodeDecode(t *testing.T) {
 	require.Equal(t, before.GetOrgName(), after.GetOrgName())
 	require.Equal(t, before.GetOrgRole(), after.GetOrgRole())
 }
+
+func TestWrapContext(t *testing.T) {
+	const key = "some-random-metadata"
+
+	ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs(key, "random-metadata"))
+	ctx, _ = identity.WithServiceIdentity(ctx, 12345)
+	var err error
+	ctx, err = wrapContext(ctx)
+	require.NoError(t, err)
+
+	outmd, ok := metadata.FromOutgoingContext(ctx)
+	require.True(t, ok)
+	val := outmd.Get(key)
+	require.Equal(t, []string{"random-metadata"}, val)
+}
+
+func TestWrapContextWithNoPreviousMetadata(t *testing.T) {
+	ctx, _ := identity.WithServiceIdentity(context.Background(), 12345)
+	ctx, err := wrapContext(ctx)
+	require.NoError(t, err)
+
+	outmd, ok := metadata.FromOutgoingContext(ctx)
+	require.True(t, ok)
+	val := outmd.Get(mdOrgID)
+	require.Equal(t, []string{"12345"}, val)
+}