Add some tests

Signed-off-by: Gabriel Adrian Samfira <[email protected]>

Author: gabriel-samfira

auth/context.go

@@ -44,8 +44,21 @@ const (
 	instanceTokenFetched contextFlags = "tokenFetched"
 	instanceHasJITConfig contextFlags = "hasJITConfig"
 	instanceParams       contextFlags = "instanceParams"
+	instanceForgeTypeKey contextFlags = "forge_type"
 )
 
+func SetInstanceForgeType(ctx context.Context, val string) context.Context {
+	return context.WithValue(ctx, instanceForgeTypeKey, val)
+}
+
+func InstanceForgeType(ctx context.Context) params.EndpointType {
+	elem := ctx.Value(instanceForgeTypeKey)
+	if elem == nil {
+		return ""
+	}
+	return elem.(params.EndpointType)
+}
+
 func SetInstanceID(ctx context.Context, id string) context.Context {
 	return context.WithValue(ctx, instanceIDKey, id)
 }
@@ -159,14 +172,15 @@ func InstanceEntity(ctx context.Context) string {
 	return elem.(string)
 }
 
-func PopulateInstanceContext(ctx context.Context, instance params.Instance) context.Context {
+func PopulateInstanceContext(ctx context.Context, instance params.Instance, claims *InstanceJWTClaims) context.Context {
 	ctx = SetInstanceID(ctx, instance.ID)
 	ctx = SetInstanceName(ctx, instance.Name)
 	ctx = SetInstancePoolID(ctx, instance.PoolID)
 	ctx = SetInstanceRunnerStatus(ctx, instance.RunnerStatus)
 	ctx = SetInstanceTokenFetched(ctx, instance.TokenFetched)
 	ctx = SetInstanceHasJITConfig(ctx, instance.JitConfiguration)
 	ctx = SetInstanceParams(ctx, instance)
+	ctx = SetInstanceForgeType(ctx, claims.ForgeType)
 	return ctx
 }
 

auth/instance_middleware.go

@@ -44,6 +44,7 @@ type InstanceJWTClaims struct {
 	// Entity is the repo or org name
 	Entity        string `json:"entity"`
 	CreateAttempt int    `json:"create_attempt"`
+	ForgeType     string `json:"forge_type"`
 	jwt.RegisteredClaims
 }
 
@@ -60,7 +61,7 @@ type instanceToken struct {
 	jwtSecret string
 }
 
-func (i *instanceToken) NewInstanceJWTToken(instance params.Instance, entity string, entityType params.ForgeEntityType, ttlMinutes uint) (string, error) {
+func (i *instanceToken) NewInstanceJWTToken(instance params.Instance, entity params.ForgeEntity, entityType params.ForgeEntityType, ttlMinutes uint) (string, error) {
 	// Token expiration is equal to the bootstrap timeout set on the pool plus the polling
 	// interval garm uses to check for timed out runners. Runners that have not sent their info
 	// by the end of this interval are most likely failed and will be reaped by garm anyway.
@@ -83,7 +84,8 @@ func (i *instanceToken) NewInstanceJWTToken(instance params.Instance, entity str
 		Name:          instance.Name,
 		PoolID:        instance.PoolID,
 		Scope:         entityType,
-		Entity:        entity,
+		Entity:        entity.String(),
+		ForgeType:     string(entity.Credentials.ForgeType),
 		CreateAttempt: instance.CreateAttempt,
 	}
 	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
@@ -124,7 +126,7 @@ func (amw *instanceMiddleware) claimsToContext(ctx context.Context, claims *Inst
 		return ctx, runnerErrors.ErrUnauthorized
 	}
 
-	ctx = PopulateInstanceContext(ctx, instanceInfo)
+	ctx = PopulateInstanceContext(ctx, instanceInfo, claims)
 	return ctx, nil
 }
 

auth/interfaces.go

@@ -26,5 +26,5 @@ type Middleware interface {
 }
 
 type InstanceTokenGetter interface {
-	NewInstanceJWTToken(instance params.Instance, entity string, poolType params.ForgeEntityType, ttlMinutes uint) (string, error)
+	NewInstanceJWTToken(instance params.Instance, entity params.ForgeEntity, poolType params.ForgeEntityType, ttlMinutes uint) (string, error)
 }