adding csa jwt generator to org resolver

Author: patrickhuie19

pkg/services/orgresolver/linking.go

@@ -8,12 +8,18 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/metadata"
 
 	log "github.com/smartcontractkit/chainlink-common/pkg/logger"
 	"github.com/smartcontractkit/chainlink-common/pkg/services"
 	linkingclient "github.com/smartcontractkit/chainlink-protos/linking-service/go/v1"
 )
 
+// JWTGenerator interface for JWT token creation
+type JWTGenerator interface {
+	CreateJWTForRequest(req any) (string, error)
+}
+
 // OrgResolver interface defines methods for resolving organization IDs from workflow owners
 type OrgResolver interface {
 	services.Service
@@ -25,6 +31,7 @@ type Config struct {
 	TLSEnabled                    bool
 	WorkflowRegistryAddress       string
 	WorkflowRegistryChainSelector uint64
+	JWTGenerator                  JWTGenerator
 }
 
 // orgResolver makes direct calls to the linking service to resolve organization IDs from workflow owners.
@@ -33,9 +40,10 @@ type orgResolver struct {
 	workflowRegistryAddress       string
 	workflowRegistryChainSelector uint64
 
-	client linkingclient.LinkingServiceClient
-	conn   *grpc.ClientConn // nil if client was injected
-	logger log.SugaredLogger
+	client       linkingclient.LinkingServiceClient
+	conn         *grpc.ClientConn // nil if client was injected
+	logger       log.SugaredLogger
+	jwtGenerator JWTGenerator
 }
 
 // NewOrgResolver creates a new org resolver with the specified configuration
@@ -49,6 +57,7 @@ func NewOrgResolverWithClient(cfg Config, client linkingclient.LinkingServiceCli
 		workflowRegistryAddress:       cfg.WorkflowRegistryAddress,
 		workflowRegistryChainSelector: cfg.WorkflowRegistryChainSelector,
 		logger:                        log.Sugared(logger).Named("OrgResolver"),
+		jwtGenerator:                  cfg.JWTGenerator,
 	}
 
 	if client != nil {
@@ -77,13 +86,36 @@ func NewOrgResolverWithClient(cfg Config, client linkingclient.LinkingServiceCli
 	return resolver, nil
 }
 
+// addJWTAuth creates and signs a JWT token, then adds it to the context
+func (o *orgResolver) addJWTAuth(ctx context.Context, req any) (context.Context, error) {
+	// Skip authentication if no JWT generator provided
+	if o.jwtGenerator == nil {
+		return ctx, nil
+	}
+
+	// Create JWT token using the JWT generator
+	jwtToken, err := o.jwtGenerator.CreateJWTForRequest(req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create JWT: %w", err)
+	}
+
+	// Add JWT to Authorization header
+	return metadata.AppendToOutgoingContext(ctx, "authorization", "Bearer "+jwtToken), nil
+}
+
 func (o *orgResolver) Get(ctx context.Context, owner string) (string, error) {
 	req := &linkingclient.GetOrganizationFromWorkflowOwnerRequest{
 		WorkflowOwner:           owner,
 		WorkflowRegistryAddress: o.workflowRegistryAddress,
 		ChainSelector:           o.workflowRegistryChainSelector,
 	}
 
+	ctx, err := o.addJWTAuth(ctx, req)
+	if err != nil {
+		o.logger.Errorw("Failed to add JWT auth to GetOrganizationFromWorkflowOwner request", "error", err)
+		return "", err
+	}
+
 	resp, err := o.client.GetOrganizationFromWorkflowOwner(ctx, req)
 	if err != nil {
 		return "", fmt.Errorf("failed to fetch organization from workflow owner: %w", err)

pkg/services/orgresolver/linking_test.go

@@ -2,12 +2,14 @@ package orgresolver
 
 import (
 	"context"
+	"errors"
 	"net"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/test/bufconn"
 
 	"github.com/smartcontractkit/chainlink-common/pkg/logger"
@@ -24,6 +26,37 @@ func (m *mockLinkingClient) GetOrganizationFromWorkflowOwner(ctx context.Context
 	}, nil
 }
 
+// mockJWTGenerator implements the JWTGenerator interface for testing
+type mockJWTGenerator struct {
+	token string
+	err   error
+}
+
+func (m *mockJWTGenerator) CreateJWTForRequest(req any) (string, error) {
+	return m.token, m.err
+}
+
+// mockLinkingClientWithAuthCheck implements the LinkingServiceClient interface and checks for authorization header
+type mockLinkingClientWithAuthCheck struct {
+	expectedAuthHeader string
+	receivedAuthHeader string
+}
+
+func (m *mockLinkingClientWithAuthCheck) GetOrganizationFromWorkflowOwner(ctx context.Context, req *linkingclient.GetOrganizationFromWorkflowOwnerRequest, opts ...grpc.CallOption) (*linkingclient.GetOrganizationFromWorkflowOwnerResponse, error) {
+	// Extract authorization header from context
+	md, ok := metadata.FromOutgoingContext(ctx)
+	if ok {
+		if authHeaders := md.Get("authorization"); len(authHeaders) > 0 {
+			m.receivedAuthHeader = authHeaders[0]
+		}
+	}
+
+	orgID := "org-" + req.WorkflowOwner
+	return &linkingclient.GetOrganizationFromWorkflowOwnerResponse{
+		OrganizationId: orgID,
+	}, nil
+}
+
 // mockLinkingServer implements the LinkingServiceServer interface for testing
 type mockLinkingServer struct {
 	linkingclient.UnimplementedLinkingServiceServer
@@ -89,7 +122,9 @@ func TestOrgResolver_NewOrgResolver_WithMockServer(t *testing.T) {
 		}),
 		grpc.WithTransportCredentials(insecure.NewCredentials()))
 	require.NoError(t, err)
-	defer conn.Close()
+	defer func() {
+		_ = conn.Close()
+	}()
 
 	client := linkingclient.NewLinkingServiceClient(conn)
 
@@ -113,3 +148,89 @@ func TestOrgResolver_NewOrgResolver_WithMockServer(t *testing.T) {
 	err = resolver.Close()
 	require.NoError(t, err)
 }
+
+func TestOrgResolver_Get_WithJWTGenerator(t *testing.T) {
+	ctx := context.Background()
+	client := &mockLinkingClientWithAuthCheck{}
+
+	// Test with JWT generator that returns a valid token
+	jwtGenerator := &mockJWTGenerator{
+		token: "test-jwt-token-123",
+		err:   nil,
+	}
+
+	cfg := Config{
+		URL:                           "test-url",
+		TLSEnabled:                    false,
+		WorkflowRegistryAddress:       "0x1234567890abcdef",
+		WorkflowRegistryChainSelector: 1,
+		JWTGenerator:                  jwtGenerator,
+	}
+
+	resolver, err := NewOrgResolverWithClient(cfg, client, logger.Test(t))
+	require.NoError(t, err)
+
+	workflowOwner := "0xabcdef1234567890"
+
+	orgID, err := resolver.Get(ctx, workflowOwner)
+	require.NoError(t, err)
+	require.Equal(t, "org-"+workflowOwner, orgID)
+
+	// Verify that the authorization header was set correctly
+	require.Equal(t, "Bearer test-jwt-token-123", client.receivedAuthHeader)
+}
+
+func TestOrgResolver_Get_WithJWTGeneratorError(t *testing.T) {
+	ctx := context.Background()
+	client := &mockLinkingClient{}
+
+	// Test with JWT generator that returns an error
+	jwtGenerator := &mockJWTGenerator{
+		token: "",
+		err:   errors.New("JWT generation failed"),
+	}
+
+	cfg := Config{
+		URL:                           "test-url",
+		TLSEnabled:                    false,
+		WorkflowRegistryAddress:       "0x1234567890abcdef",
+		WorkflowRegistryChainSelector: 1,
+		JWTGenerator:                  jwtGenerator,
+	}
+
+	resolver, err := NewOrgResolverWithClient(cfg, client, logger.Test(t))
+	require.NoError(t, err)
+
+	workflowOwner := "0xabcdef1234567890"
+
+	// The Get call should fail due to JWT generation error
+	_, err = resolver.Get(ctx, workflowOwner)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "JWT generation failed")
+}
+
+func TestOrgResolver_Get_WithoutJWTGenerator(t *testing.T) {
+	ctx := context.Background()
+	client := &mockLinkingClientWithAuthCheck{}
+
+	// Test without JWT generator (should not set authorization header)
+	cfg := Config{
+		URL:                           "test-url",
+		TLSEnabled:                    false,
+		WorkflowRegistryAddress:       "0x1234567890abcdef",
+		WorkflowRegistryChainSelector: 1,
+		JWTGenerator:                  nil, // No JWT generator
+	}
+
+	resolver, err := NewOrgResolverWithClient(cfg, client, logger.Test(t))
+	require.NoError(t, err)
+
+	workflowOwner := "0xabcdef1234567890"
+
+	orgID, err := resolver.Get(ctx, workflowOwner)
+	require.NoError(t, err)
+	require.Equal(t, "org-"+workflowOwner, orgID)
+
+	// Verify that no authorization header was set
+	require.Empty(t, client.receivedAuthHeader)
+}