diff --git a/pkg/services/orgresolver/linking.go b/pkg/services/orgresolver/linking.go index 3b05d52d9..40be90ce4 100644 --- a/pkg/services/orgresolver/linking.go +++ b/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,6 +86,23 @@ 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, @@ -84,6 +110,12 @@ func (o *orgResolver) Get(ctx context.Context, owner string) (string, error) { 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) diff --git a/pkg/services/orgresolver/linking_test.go b/pkg/services/orgresolver/linking_test.go index 8b0a9836f..f1c2b3efb 100644 --- a/pkg/services/orgresolver/linking_test.go +++ b/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) +}