feat: load membership from header (#1654)

Signed-off-by: Miguel Martinez <[email protected]>

Author: migmartri

app/controlplane/internal/usercontext/currentorganization_middleware.go

@@ -36,10 +36,23 @@ func WithCurrentOrganizationMiddleware(userUseCase biz.UserOrgFinder, logger *lo
 				return handler(ctx, req)
 			}
 
-			var err error
-			ctx, err = setCurrentOrganization(ctx, u, userUseCase, logger)
+			orgName, err := getOrganizationName(ctx)
 			if err != nil {
-				return nil, fmt.Errorf("error setting current org: %w", err)
+				return nil, fmt.Errorf("error getting organization name: %w", err)
+			}
+
+			if orgName != "" {
+				ctx, err = setCurrentOrganizationFromHeader(ctx, u, orgName, userUseCase)
+				if err != nil {
+					return nil, fmt.Errorf("error setting current org: %w", err)
+				}
+			} else {
+				// If no organization name is provided, we use the DB to find the current organization
+				// DEPRECATED: in favor of header based org selection
+				ctx, err = setCurrentOrganizationFromDB(ctx, u, userUseCase, logger)
+				if err != nil {
+					return nil, fmt.Errorf("error setting current org: %w", err)
+				}
 			}
 
 			org := entities.CurrentOrg(ctx)
@@ -54,8 +67,19 @@ func WithCurrentOrganizationMiddleware(userUseCase biz.UserOrgFinder, logger *lo
 	}
 }
 
+func setCurrentOrganizationFromHeader(ctx context.Context, user *entities.User, orgName string, userUC biz.UserOrgFinder) (context.Context, error) {
+	membership, err := userUC.MembershipInOrg(ctx, user.ID, orgName)
+	if err != nil {
+		return nil, fmt.Errorf("failed to find membership: %w", err)
+	}
+
+	ctx = entities.WithCurrentOrg(ctx, &entities.Org{Name: membership.Org.Name, ID: membership.Org.ID, CreatedAt: membership.CreatedAt})
+	// Set the authorization subject that will be used to check the policies
+	return WithAuthzSubject(ctx, string(membership.Role)), nil
+}
+
 // Find the current membership of the user and sets it on the context
-func setCurrentOrganization(ctx context.Context, user *entities.User, userUC biz.UserOrgFinder, logger *log.Helper) (context.Context, error) {
+func setCurrentOrganizationFromDB(ctx context.Context, user *entities.User, userUC biz.UserOrgFinder, logger *log.Helper) (context.Context, error) {
 	// We load the current organization
 	membership, err := userUC.CurrentMembership(ctx, user.ID)
 	if err != nil {

app/controlplane/internal/usercontext/usercontext.go

@@ -36,3 +36,13 @@ func GetRawToken(ctx context.Context) (string, error) {
 	}
 	return auths[1], nil
 }
+
+func getOrganizationName(ctx context.Context) (string, error) {
+	const OrganizationHeader = "Chainloop-Organization"
+	header, ok := transport.FromServerContext(ctx)
+	if ok {
+		return header.RequestHeader().Get(OrganizationHeader), nil
+	}
+
+	return "", nil
+}

app/controlplane/pkg/biz/mocks/UserOrgFinder.go

@@ -46,6 +46,7 @@ type UserRepo interface {
 type UserOrgFinder interface {
 	FindByID(ctx context.Context, userID string) (*User, error)
 	CurrentMembership(ctx context.Context, userID string) (*Membership, error)
+	MembershipInOrg(ctx context.Context, userID string, orgName string) (*Membership, error)
 }
 
 type UserUseCase struct {
@@ -159,6 +160,17 @@ func (uc *UserUseCase) FindByID(ctx context.Context, userID string) (*User, erro
 	return uc.userRepo.FindByID(ctx, userUUID)
 }
 
+func (uc *UserUseCase) MembershipInOrg(ctx context.Context, userID string, orgName string) (*Membership, error) {
+	m, err := uc.membershipUseCase.FindByOrgNameAndUser(ctx, orgName, userID)
+	if err != nil {
+		return nil, fmt.Errorf("failed to find membership: %w", err)
+	} else if m == nil {
+		return nil, NewErrNotFound("user does not have this org associated")
+	}
+
+	return m, nil
+}
+
 // Find the membership associated with the user that's marked as current
 // If none is selected, it will pick the first one and set it as current
 func (uc *UserUseCase) CurrentMembership(ctx context.Context, userID string) (*Membership, error) {