diff --git a/pkg/grpcutil/grpcutil.go b/pkg/grpcutil/grpcutil.go new file mode 100644 index 000000000..9dd664a67 --- /dev/null +++ b/pkg/grpcutil/grpcutil.go @@ -0,0 +1,31 @@ +// Package grpcutil contains functions for interacting with gRPC. +package grpcutil + +import ( + "context" + + "google.golang.org/grpc/metadata" +) + +// SessionIDMetadataKey is the key in the metadata. +const SessionIDMetadataKey = "sessionid" + +// WithOutgoingSessionID appends a metadata header for the session ID to a context. +func WithOutgoingSessionID(ctx context.Context, sessionID string) context.Context { + return metadata.AppendToOutgoingContext(ctx, SessionIDMetadataKey, sessionID) +} + +// SessionIDFromGRPCRequest returns the session id from the gRPC request. +func SessionIDFromGRPCRequest(ctx context.Context) (sessionID string, ok bool) { + md, ok := metadata.FromIncomingContext(ctx) + if !ok { + return "", false + } + + sessionIDs := md.Get(SessionIDMetadataKey) + if len(sessionIDs) == 0 { + return "", false + } + + return sessionIDs[0], true +} diff --git a/pkg/grpcutil/grpcutil_test.go b/pkg/grpcutil/grpcutil_test.go new file mode 100644 index 000000000..ce076cfd0 --- /dev/null +++ b/pkg/grpcutil/grpcutil_test.go @@ -0,0 +1,29 @@ +package grpcutil + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "google.golang.org/grpc/metadata" +) + +func TestWithOutgoingSessionID(t *testing.T) { + ctx := context.Background() + ctx = WithOutgoingSessionID(ctx, "EXAMPLE") + md, ok := metadata.FromOutgoingContext(ctx) + if !assert.True(t, ok) { + return + } + assert.Equal(t, []string{"EXAMPLE"}, md.Get("sessionid")) +} + +func TestSessionIDFromGRPCRequest(t *testing.T) { + ctx := context.Background() + ctx = metadata.NewIncomingContext(ctx, metadata.MD{ + "sessionid": {"EXAMPLE"}, + }) + sessionID, ok := SessionIDFromGRPCRequest(ctx) + assert.True(t, ok) + assert.Equal(t, "EXAMPLE", sessionID) +}