feat(attestations): support API tokens for attestations (#763)

Signed-off-by: Javier Rodriguez <[email protected]>

Author: javirln

app/controlplane/api/controlplane/v1/workflow_run.pb.go

@@ -43,6 +43,9 @@ service WorkflowRunService {
 
 message AttestationServiceGetContractRequest {
   int32 contract_revision = 1;
+  // This parameter is not needed by Robot Account since they have the workflowID embedded.
+  // API Tokens will send the parameter explicitly
+  string workflow_name = 2;
 }
 
 message AttestationServiceGetContractResponse {

app/controlplane/api/controlplane/v1/workflow_run.proto

@@ -203,10 +203,11 @@ func (r *WorkflowRepo) FindByID(ctx context.Context, id uuid.UUID) (*biz.Workflo
 		Where(workflow.DeletedAtIsNil(), workflow.ID(id)).
 		WithContract().WithOrganization().
 		Only(ctx)
-	if err != nil && !ent.IsNotFound(err) {
+	if err != nil {
+		if ent.IsNotFound(err) {
+			return nil, biz.NewErrNotFound("workflow")
+		}
 		return nil, err
-	} else if workflow == nil {
-		return nil, nil
 	}
 
 	// Not efficient, we need to do a query limit = 1 grouped by workflowID

app/controlplane/api/gen/frontend/controlplane/v1/workflow_run.ts

@@ -26,17 +26,17 @@ import (
 	authzMiddleware "github.com/chainloop-dev/chainloop/app/controlplane/internal/authz/middleware"
 	"github.com/chainloop-dev/chainloop/app/controlplane/internal/biz"
 	"github.com/chainloop-dev/chainloop/app/controlplane/internal/conf"
-	"github.com/chainloop-dev/chainloop/app/controlplane/internal/jwt/robotaccount"
 	"github.com/chainloop-dev/chainloop/app/controlplane/internal/jwt/user"
+	"github.com/chainloop-dev/chainloop/app/controlplane/internal/usercontext/attjwtmiddleware"
 
 	"github.com/bufbuild/protovalidate-go"
 	"github.com/chainloop-dev/chainloop/app/controlplane/internal/service"
 	"github.com/chainloop-dev/chainloop/app/controlplane/internal/usercontext"
 	"github.com/chainloop-dev/chainloop/internal/credentials"
 	"github.com/getsentry/sentry-go"
-	jwt "github.com/golang-jwt/jwt/v4"
+	"github.com/golang-jwt/jwt/v4"
 
-	errors "github.com/go-kratos/kratos/v2/errors"
+	"github.com/go-kratos/kratos/v2/errors"
 	"github.com/go-kratos/kratos/v2/log"
 	"github.com/go-kratos/kratos/v2/middleware"
 	jwtMiddleware "github.com/go-kratos/kratos/v2/middleware/auth/jwt"
@@ -58,6 +58,7 @@ type Opts struct {
 	ReferrerUseCase      *biz.ReferrerUseCase
 	APITokenUseCase      *biz.APITokenUseCase
 	OrganizationUserCase *biz.OrganizationUseCase
+	WorkflowUseCase      *biz.WorkflowUseCase
 	// Services
 	WorkflowSvc         *service.WorkflowService
 	AuthSvc             *service.AuthService
@@ -192,16 +193,17 @@ func craftMiddleware(opts *Opts) []middleware.Middleware {
 	middlewares = append(middlewares,
 		// if we require a robot account
 		selector.Server(
-			// 1 - Extract the robot account from the JWT
-			jwtMiddleware.Server(func(_ *jwt.Token) (interface{}, error) {
-				// TODO: add support to multiple signing methods and keys
-				return []byte(opts.AuthConfig.GeneratedJwsHmacSecret), nil
-			},
-				jwtMiddleware.WithSigningMethod(robotaccount.SigningMethod),
-				jwtMiddleware.WithClaims(func() jwt.Claims { return &robotaccount.CustomClaims{} }),
+			// 1 - Extract information from the JWT by using the claims
+			attjwtmiddleware.WithJWTMulti(
+				// Robot account provider
+				attjwtmiddleware.NewRobotAccountProvider(opts.AuthConfig.GeneratedJwsHmacSecret),
+				// API Token provider
+				attjwtmiddleware.NewAPITokenProvider(opts.AuthConfig.GeneratedJwsHmacSecret),
 			),
-			// 2 - Set its workflow and organization in the context
-			usercontext.WithCurrentRobotAccount(opts.RobotAccountUseCase, logHelper),
+			// 2.a - Set its workflow and organization in the context
+			usercontext.WithAttestationContextFromRobotAccount(opts.RobotAccountUseCase, logHelper),
+			// 2.b - Set its API token and Robot Account as alternative to the user
+			usercontext.WithAttestationContextFromAPIToken(opts.APITokenUseCase, logHelper),
 		).Match(requireRobotAccountMatcher()).Build(),
 	)
 

app/controlplane/cmd/wire_gen.go

@@ -90,7 +90,22 @@ func NewAttestationService(opts *NewAttestationServiceOpts) *AttestationService
 func (s *AttestationService) GetContract(ctx context.Context, req *cpAPI.AttestationServiceGetContractRequest) (*cpAPI.AttestationServiceGetContractResponse, error) {
 	robotAccount := usercontext.CurrentRobotAccount(ctx)
 	if robotAccount == nil {
-		return nil, errors.NotFound("not found", "robot account not found")
+		return nil, errors.NotFound("not found", "neither robot account nor API token found")
+	}
+
+	// If WorkflowID is not set and no workflowName is provided, return BadRequest.
+	if req.GetWorkflowName() == "" && robotAccount.WorkflowID == "" {
+		return nil, errors.BadRequest("bad request", "when using an API Token, workflow name is required as parameter")
+	}
+
+	// If WorkflowID is not set, retrieve it using the workflow name. This is needed since when coming from
+	// an API Token request, the only information we have set is the workflow name not the ID
+	if robotAccount.WorkflowID == "" {
+		workflow, err := s.workflowUseCase.FindByNameInOrg(ctx, robotAccount.OrgID, req.GetWorkflowName())
+		if err != nil {
+			return nil, fmt.Errorf("error retrieving the workflow: %w", err)
+		}
+		robotAccount.WorkflowID = workflow.ID.String()
 	}
 
 	// Find workflow

app/controlplane/internal/data/workflow.go

@@ -27,6 +27,7 @@ import (
 	"github.com/chainloop-dev/chainloop/app/controlplane/internal/authz"
 	"github.com/chainloop-dev/chainloop/app/controlplane/internal/biz"
 	"github.com/chainloop-dev/chainloop/app/controlplane/internal/jwt/apitoken"
+	"github.com/chainloop-dev/chainloop/app/controlplane/internal/usercontext/attjwtmiddleware"
 	"github.com/go-kratos/kratos/v2/middleware"
 	jwtMiddleware "github.com/go-kratos/kratos/v2/middleware/auth/jwt"
 )
@@ -101,6 +102,74 @@ func WithCurrentAPITokenAndOrgMiddleware(apiTokenUC *biz.APITokenUseCase, orgUC
 	}
 }
 
+// WithAttestationContextFromAPIToken injects the API-Token, organization + robot account to the context
+func WithAttestationContextFromAPIToken(apiTokenUC *biz.APITokenUseCase, logger *log.Helper) middleware.Middleware {
+	return func(handler middleware.Handler) middleware.Handler {
+		return func(ctx context.Context, req interface{}) (interface{}, error) {
+			authInfo, ok := attjwtmiddleware.FromJWTAuthContext(ctx)
+			// If not found means that there is no currentUser set in the context
+			if !ok {
+				logger.Warn("couldn't extract org/user, JWT parser middleware not running before this one?")
+				return nil, errors.New("can't extract JWT info from the context")
+			}
+
+			// If the token is not an API token, we don't need to do anything
+			if authInfo.ProviderKey != attjwtmiddleware.APITokenProviderKey {
+				return handler(ctx, req)
+			}
+
+			genericClaims, ok := authInfo.Claims.(jwt.MapClaims)
+			if !ok {
+				return nil, errors.New("error mapping the claims")
+			}
+
+			// We've received an API-token, double check its audience
+			if !genericClaims.VerifyAudience(apitoken.Audience, true) {
+				return nil, errors.New("unexpected token, invalid audience")
+			}
+
+			var err error
+			tokenID, ok := genericClaims["jti"].(string)
+			if !ok || tokenID == "" {
+				return nil, errors.New("error mapping the API-token claims")
+			}
+
+			ctx, err = setRobotAccountFromAPIToken(ctx, apiTokenUC, tokenID)
+			if err != nil {
+				return nil, errors.New("error extracting organization from APIToken")
+			}
+
+			logger.Infow("msg", "[authN] processed credentials", "id", tokenID, "type", "API-token")
+
+			return handler(ctx, req)
+		}
+	}
+}
+
+func setRobotAccountFromAPIToken(ctx context.Context, apiTokenUC *biz.APITokenUseCase, tokenID string) (context.Context, error) {
+	if tokenID == "" {
+		return nil, errors.New("error retrieving the key ID from the API token")
+	}
+
+	// Check that the token exists and is not revoked
+	token, err := apiTokenUC.FindByID(ctx, tokenID)
+	if err != nil {
+		return nil, fmt.Errorf("error retrieving the API token: %w", err)
+	} else if token == nil {
+		return nil, errors.New("API token not found")
+	}
+
+	// Note: Expiration time does not need to be checked because that's done at the JWT
+	// verification layer, which happens before this middleware is called
+	if token.RevokedAt != nil {
+		return nil, errors.New("API token revoked")
+	}
+
+	ctx = withRobotAccount(ctx, &RobotAccount{OrgID: token.OrganizationID.String()})
+
+	return ctx, nil
+}
+
 // Set the current organization and API-Token in the context
 func setCurrentOrgAndAPIToken(ctx context.Context, apiTokenUC *biz.APITokenUseCase, orgUC *biz.OrganizationUseCase, tokenID string) (context.Context, error) {
 	if tokenID == "" {