From 9031a4029e257976f7719ae250fb806dadb0124c Mon Sep 17 00:00:00 2001
From: Gabriel Adrian Samfira <gsamfira@cloudbasesolutions.com>
Date: Tue, 30 Jan 2024 11:07:55 +0000
Subject: [PATCH] More strict instance token checks

This change invalidates tokens based on more parameters. Tokens that were
generated for previous attempts of spinning up an instance will be invalidates.

Also, only instances that are in Running or Creating will be able to authenticate.

Signed-off-by: Gabriel Adrian Samfira <gsamfira@cloudbasesolutions.com>
---
 auth/instance_middleware.go | 48 +++++++++++++++++++++++++++++++------
 1 file changed, 41 insertions(+), 7 deletions(-)

diff --git a/auth/instance_middleware.go b/auth/instance_middleware.go
index 71147c78..1562b8b9 100644
--- a/auth/instance_middleware.go
+++ b/auth/instance_middleware.go
@@ -17,16 +17,18 @@ package auth
 import (
 	"context"
 	"fmt"
+	"log/slog"
 	"net/http"
 	"strings"
 	"time"
 
-	runnerErrors "github.com/cloudbase/garm-provider-common/errors"
 	"github.com/cloudbase/garm/config"
 	dbCommon "github.com/cloudbase/garm/database/common"
 	"github.com/cloudbase/garm/params"
 	"github.com/cloudbase/garm/runner/common"
 
+	runnerErrors "github.com/cloudbase/garm-provider-common/errors"
+	commonParams "github.com/cloudbase/garm-provider-common/params"
 	jwt "github.com/golang-jwt/jwt/v5"
 	"github.com/pkg/errors"
 )
@@ -39,7 +41,8 @@ type InstanceJWTClaims struct {
 	// Scope is either repository or organization
 	Scope params.PoolType `json:"scope"`
 	// Entity is the repo or org name
-	Entity string `json:"entity"`
+	Entity        string `json:"entity"`
+	CreateAttempt int    `json:"create_attempt"`
 	jwt.RegisteredClaims
 }
 
@@ -56,11 +59,12 @@ func NewInstanceJWTToken(instance params.Instance, secret, entity string, poolTy
 			ExpiresAt: expires,
 			Issuer:    "garm",
 		},
-		ID:     instance.ID,
-		Name:   instance.Name,
-		PoolID: instance.PoolID,
-		Scope:  poolType,
-		Entity: entity,
+		ID:            instance.ID,
+		Name:          instance.Name,
+		PoolID:        instance.PoolID,
+		Scope:         poolType,
+		Entity:        entity,
+		CreateAttempt: instance.CreateAttempt,
 	}
 	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
 	tokenString, err := token.SignedString([]byte(secret))
@@ -157,6 +161,36 @@ func (amw *instanceMiddleware) Middleware(next http.Handler) http.Handler {
 			return
 		}
 
+		instanceParams, err := InstanceParams(ctx)
+		if err != nil {
+			slog.InfoContext(
+				ctx, "could not find instance params",
+				"runner_name", InstanceName(ctx))
+			invalidAuthResponse(ctx, w)
+			return
+		}
+
+		// Token was generated for a previous attempt at creating this instance.
+		if claims.CreateAttempt != instanceParams.CreateAttempt {
+			slog.InfoContext(
+				ctx, "invalid token create attempt",
+				"runner_name", InstanceName(ctx),
+				"token_create_attempt", claims.CreateAttempt,
+				"instance_create_attempt", instanceParams.CreateAttempt)
+			invalidAuthResponse(ctx, w)
+			return
+		}
+
+		// Only allow instances that are in the creating or running state to authenticate.
+		if instanceParams.Status != commonParams.InstanceCreating && instanceParams.Status != commonParams.InstanceRunning {
+			slog.InfoContext(
+				ctx, "invalid instance status",
+				"runner_name", InstanceName(ctx),
+				"status", instanceParams.Status)
+			invalidAuthResponse(ctx, w)
+			return
+		}
+
 		next.ServeHTTP(w, r.WithContext(ctx))
 	})
 }