fix: ECS RunTask and CreateLaunchTemplateVersion honour ClientToken (#974)

* fix: ECS RunTask and CreateLaunchTemplateVersion honour ClientToken

A retried RunTask placed the tasks again: it mints a fresh task ID per task,
reserves CPU/memory/GPU, hot-plugs an ENI in awsvpc mode and starts the
containers, so a resend with Count 5 left ten running tasks, double the reserved
capacity and five leaked ENIs. CreateLaunchTemplateVersion added a second
identical version and moved $Latest.

RunTask gets a per-action wrapper over idempotency.Do rather than a dispatch
decorator: it is the only ECS action that needs one. A partial placement
finalizes rather than aborts, since re-running would add tasks on top of the
ones already started.

CreateLaunchTemplateVersion reuses the EC2 dispatch decorator, so it is an
opt-in rather than new machinery.

CreateLaunchTemplate, ModifyLaunchTemplate and ECS CreateService accept a token
they do not store; each is name-keyed or a CAS and cannot duplicate, which is
now recorded at the handler.

* test: cover the RunTask client-token binding layer

Split the token layer out of RunTask's NATS binding as runTaskIdempotent, the
same seam runTask already has, so the no-token, store-failure and replay paths
are reachable without a live connection.
This commit is contained in:
Josh Siviour
2026-09-01 21:54:32 +10:00
committed by GitHub
parent 44abc812b6
commit 687974995f
7 changed files with 425 additions and 4 deletions
+1 -1
View File
@@ -412,7 +412,7 @@ var ec2Actions = map[string]ec2Action{
"CreateLaunchTemplate": ec2Handler(func(ctx context.Context, input *ec2.CreateLaunchTemplateInput, gw *GatewayConfig, accountID string) (any, error) {
return gateway_ec2_launchtemplate.CreateLaunchTemplate(ctx, input, gw.NATSConn, accountID)
}),
"CreateLaunchTemplateVersion": ec2Handler(func(ctx context.Context, input *ec2.CreateLaunchTemplateVersionInput, gw *GatewayConfig, accountID string) (any, error) {
"CreateLaunchTemplateVersion": ec2IdempotentHandler(func(ctx context.Context, input *ec2.CreateLaunchTemplateVersionInput, gw *GatewayConfig, accountID string) (ec2.CreateLaunchTemplateVersionOutput, error) {
return gateway_ec2_launchtemplate.CreateLaunchTemplateVersion(ctx, input, gw.NATSConn, accountID)
}),
"DeleteLaunchTemplate": ec2Handler(func(ctx context.Context, input *ec2.DeleteLaunchTemplateInput, gw *GatewayConfig, accountID string) (any, error) {
+19
View File
@@ -185,6 +185,24 @@ func TestIdempotentHandler_SameTokenCreatesOneResource(t *testing.T) {
assert.Contains(t, string(second), "ami-1")
})
// A duplicate would add a second identical version and move $Latest under a
// client that thought it made one call.
t.Run("CreateLaunchTemplateVersion", func(t *testing.T) {
t.Parallel()
gw := newTokenGateway(t)
calls, first, second := dispatchTwice(t, gw, "CreateLaunchTemplateVersion",
&ec2.CreateLaunchTemplateVersionInput{LaunchTemplateId: aws.String("lt-1"), ClientToken: tok},
ec2.CreateLaunchTemplateVersionOutput{
LaunchTemplateVersion: &ec2.LaunchTemplateVersion{
LaunchTemplateId: aws.String("lt-1"),
VersionNumber: aws.Int64(2),
},
})
assert.Equal(t, 1, calls)
assert.Equal(t, responseFields(t, first), responseFields(t, second))
assert.Contains(t, string(second), "<versionNumber>2</versionNumber>")
})
t.Run("CreateCapacityReservation", func(t *testing.T) {
t.Parallel()
gw := newTokenGateway(t)
@@ -288,6 +306,7 @@ func TestEC2Actions_CreatesHonourClientToken(t *testing.T) {
"CopyImage",
"CreateCapacityReservation",
"CreateEgressOnlyInternetGateway",
"CreateLaunchTemplateVersion",
"CreateNatGateway",
"CreateNetworkInterface",
"CreateRouteTable",
+54 -2
View File
@@ -3,9 +3,12 @@ package gateway_ecs
import (
"context"
"encoding/json"
"errors"
"log/slog"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ecs"
"github.com/mulgadc/spinifex/spinifex/awserrors"
handlers_ecs "github.com/mulgadc/spinifex/spinifex/handlers/ecs"
"github.com/nats-io/nats.go"
)
@@ -177,17 +180,66 @@ func ListContainerInstances(ctx context.Context, nc *nats.Conn, accountID string
// --- Task ---
// RunTask honours ClientToken: a placement reserves capacity, hot-plugs an ENI
// in awsvpc mode and starts containers, so a retried request would double all
// three rather than leave one stray resource behind.
func RunTask(ctx context.Context, nc *nats.Conn, accountID string, body []byte, passRoleCheck PassRoleChecker) (any, error) {
input := new(ecs.RunTaskInput)
if err := unmarshalIfBody(body, input); err != nil {
return nil, err
}
return runTask(ctx, handlers_ecs.NewNATSECSService(nc), accountID, input, passRoleCheck)
out, err := runTaskIdempotent(ctx, handlers_ecs.NewNATSECSService(nc), accountID, input, passRoleCheck,
func() (*runTaskStore, error) { return getRunTaskStore(ctx, nc) })
if err != nil {
// Returned as a typed nil otherwise, which is a non-nil any.
return nil, err
}
return out, nil
}
// runTaskIdempotent is RunTask's token layer, split from the NATS binding the
// way runTask is. openStore is a function so an untokened request never pays
// the bind, and so a test can supply a store without a live connection.
func runTaskIdempotent(
ctx context.Context,
svc handlers_ecs.ECSService,
accountID string,
input *ecs.RunTaskInput,
passRoleCheck PassRoleChecker,
openStore func() (*runTaskStore, error),
) (*ecs.RunTaskOutput, error) {
// No token means the caller did not ask for idempotency, as on AWS.
token := aws.StringValue(input.ClientToken)
if token == "" {
return runTask(ctx, svc, accountID, input, passRoleCheck)
}
store, serr := openStore()
if serr != nil {
// Launching anyway would place the duplicate tasks the token was sent to
// prevent, so this fails rather than degrading to no idempotency.
slog.ErrorContext(ctx, "ECS RunTask: client-token store unavailable", "err", serr)
return nil, errors.New(awserrors.ErrorServerInternal)
}
// Hashed before the launch, so a handler that mutates its input cannot change
// what a retry of the same token hashes to.
paramHash := runTaskParamHash(input)
out, err := runTaskWithClientToken(ctx, store, accountID, token, paramHash, func() (ecs.RunTaskOutput, error) {
res, rerr := runTask(ctx, svc, accountID, input, passRoleCheck)
if rerr != nil {
return ecs.RunTaskOutput{}, rerr
}
return *res, nil
})
if err != nil {
return nil, err
}
return &out, nil
}
// runTask is RunTask's core, split out so tests can inject a fake ECSService
// instead of a live NATS connection.
func runTask(ctx context.Context, svc handlers_ecs.ECSService, accountID string, input *ecs.RunTaskInput, passRoleCheck PassRoleChecker) (any, error) {
func runTask(ctx context.Context, svc handlers_ecs.ECSService, accountID string, input *ecs.RunTaskInput, passRoleCheck PassRoleChecker) (*ecs.RunTaskOutput, error) {
if err := checkTaskLaunchRoles(ctx, svc, accountID, input.TaskDefinition, passRoleCheck); err != nil {
return nil, err
}
+86
View File
@@ -0,0 +1,86 @@
package gateway_ecs
import (
"context"
"errors"
"fmt"
"log/slog"
"sync"
"time"
"github.com/aws/aws-sdk-go/service/ecs"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/mulgadc/spinifex/spinifex/idempotency"
"github.com/nats-io/nats.go"
"github.com/nats-io/nats.go/jetstream"
)
const (
// kvBucketClientTokens is the JetStream KV bucket for ECS ClientToken
// records. ECS keeps its own rather than borrowing EC2's, whose name and TTL
// belong to the EC2 gateway.
kvBucketClientTokens = "spinifex-ecs-clienttokens" //nolint:gosec // G101 false positive: KV bucket name, not a credential
// clientTokenTTL must outlast SDK retry windows; short enough that a crashed
// in-flight record ages out and frees the token for a fresh attempt.
clientTokenTTL = 15 * time.Minute
)
// runTaskStore replays a whole RunTaskOutput: a placement can partly succeed, so
// the failures are as much a part of the result as the tasks.
type runTaskStore = idempotency.Store[ecs.RunTaskOutput]
var (
ctStore *runTaskStore
ctOnce sync.Once
errCTStoreInit error
)
// getRunTaskStore lazily binds the process-wide store. Process-wide because the
// ECS handlers are free functions taking a *nats.Conn, with no instance to hang
// it on.
func getRunTaskStore(ctx context.Context, nc *nats.Conn) (*runTaskStore, error) {
ctOnce.Do(func() {
js, err := jetstream.New(nc)
if err != nil {
errCTStoreInit = fmt.Errorf("clienttoken jetstream: %w", err)
return
}
// The bind happens once per process, so it must not inherit the first
// caller's cancellation: a client that disconnects mid-open would poison
// the store for every later launch.
ctStore, errCTStoreInit = idempotency.OpenStore[ecs.RunTaskOutput](
context.WithoutCancel(ctx), js, kvBucketClientTokens, clientTokenTTL)
})
return ctStore, errCTStoreInit
}
// runTaskParamHash hashes the request excluding ClientToken, so the same params
// always produce the same hash. Must run before any input mutation.
func runTaskParamHash(input *ecs.RunTaskInput) string {
clone := *input
clone.ClientToken = nil
return idempotency.ParamHash(&clone)
}
// runTaskWithClientToken wraps a launch in ClientToken idempotency and maps the
// token errors onto the AWS ones RunTask returns. Extracted for unit-testability.
//
// A partial placement finalizes rather than aborts: some tasks are already
// running, and re-running would add more on top of them.
func runTaskWithClientToken(
ctx context.Context,
store *runTaskStore,
accountID, token, paramHash string,
launch func() (ecs.RunTaskOutput, error),
) (ecs.RunTaskOutput, error) {
out, err := idempotency.Do(ctx, store, accountID, token, paramHash, launch)
switch {
case errors.Is(err, idempotency.ErrParamMismatch):
return ecs.RunTaskOutput{}, errors.New(awserrors.ErrorIdempotentParameterMismatch)
case errors.Is(err, idempotency.ErrUnavailable):
slog.ErrorContext(ctx, "ECS RunTask: client-token claim failed", "token", token, "err", err)
return ecs.RunTaskOutput{}, errors.New(awserrors.ErrorServerInternal)
}
return out, err
}
+258
View File
@@ -0,0 +1,258 @@
//test:in-package — runTaskWithClientToken and runTaskParamHash are unexported,
//and the rest of this package's tests are in-package for the same reason.
package gateway_ecs
import (
"errors"
"testing"
"time"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ecs"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/mulgadc/spinifex/spinifex/idempotency"
"github.com/mulgadc/spinifex/spinifex/testutil"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
const tokenTestAccount = "123456789012"
func newRunTaskStore(t *testing.T) *runTaskStore {
t.Helper()
_, nc, _ := testutil.StartTestJetStream(t)
store, err := idempotency.OpenStore[ecs.RunTaskOutput](
t.Context(), testutil.NewJetStream(t, nc), "ecs-tokens", time.Minute)
require.NoError(t, err)
return store
}
func taskOutput(ids ...string) ecs.RunTaskOutput {
out := ecs.RunTaskOutput{}
for _, id := range ids {
out.Tasks = append(out.Tasks, &ecs.Task{TaskArn: aws.String(id)})
}
return out
}
// The bead's criterion for RunTask: one token, two calls, one placement. A
// second placement would double the reserved capacity and leak an ENI per task.
func TestRunTaskWithClientToken_PlacesOnce(t *testing.T) {
t.Parallel()
store := newRunTaskStore(t)
const tok, hash = "tok-1", "h"
launches := 0
launch := func() (ecs.RunTaskOutput, error) {
launches++
return taskOutput("task-a", "task-b"), nil
}
first, err := runTaskWithClientToken(t.Context(), store, tokenTestAccount, tok, hash, launch)
require.NoError(t, err)
second, err := runTaskWithClientToken(t.Context(), store, tokenTestAccount, tok, hash, launch)
require.NoError(t, err)
assert.Equal(t, 1, launches, "the duplicate must replay, not place again")
require.Len(t, second.Tasks, 2)
assert.Equal(t, first.Tasks[0].TaskArn, second.Tasks[0].TaskArn)
assert.Equal(t, first.Tasks[1].TaskArn, second.Tasks[1].TaskArn)
}
// A placement where some tasks started and others failed is still a placement:
// re-running would add tasks on top of the ones already running.
func TestRunTaskWithClientToken_PartialSuccessReplays(t *testing.T) {
t.Parallel()
store := newRunTaskStore(t)
const tok, hash = "tok-partial", "h"
partial := taskOutput("task-a")
partial.Failures = []*ecs.Failure{{Reason: aws.String("RESOURCE:placement")}}
launches := 0
launch := func() (ecs.RunTaskOutput, error) {
launches++
return partial, nil
}
_, err := runTaskWithClientToken(t.Context(), store, tokenTestAccount, tok, hash, launch)
require.NoError(t, err)
replay, err := runTaskWithClientToken(t.Context(), store, tokenTestAccount, tok, hash, launch)
require.NoError(t, err)
assert.Equal(t, 1, launches, "a partial success must not be retried under the same token")
require.Len(t, replay.Tasks, 1)
require.Len(t, replay.Failures, 1, "the failures replay alongside the tasks")
assert.Equal(t, "RESOURCE:placement", aws.StringValue(replay.Failures[0].Reason))
}
// A launch that errored placed nothing, so the token has to free up for a real
// retry rather than replay a result that never existed.
func TestRunTaskWithClientToken_FailedLaunchIsRetryable(t *testing.T) {
t.Parallel()
store := newRunTaskStore(t)
const tok, hash = "tok-2", "h"
boom := errors.New("placement failed")
_, err := runTaskWithClientToken(t.Context(), store, tokenTestAccount, tok, hash,
func() (ecs.RunTaskOutput, error) { return ecs.RunTaskOutput{}, boom })
require.ErrorIs(t, err, boom, "the launch's own error reaches the caller unchanged")
launches := 0
out, err := runTaskWithClientToken(t.Context(), store, tokenTestAccount, tok, hash,
func() (ecs.RunTaskOutput, error) {
launches++
return taskOutput("task-a"), nil
})
require.NoError(t, err)
assert.Equal(t, 1, launches)
require.Len(t, out.Tasks, 1)
}
// Reusing a token with different parameters is the AWS mismatch case, and must
// not place anything.
func TestRunTaskWithClientToken_ParamMismatchIsRejected(t *testing.T) {
t.Parallel()
store := newRunTaskStore(t)
const tok = "tok-3"
_, err := runTaskWithClientToken(t.Context(), store, tokenTestAccount, tok, "hash-a",
func() (ecs.RunTaskOutput, error) { return taskOutput("task-a"), nil })
require.NoError(t, err)
launched := false
_, err = runTaskWithClientToken(t.Context(), store, tokenTestAccount, tok, "hash-b",
func() (ecs.RunTaskOutput, error) {
launched = true
return ecs.RunTaskOutput{}, nil
})
require.Error(t, err)
assert.Equal(t, awserrors.ErrorIdempotentParameterMismatch, err.Error())
assert.False(t, launched, "a mismatched token must not place tasks")
}
// The hash must ignore the token: a retry carries the same token and params, and
// changing Count is a different request even under the same token.
func TestRunTaskParamHash_IgnoresTheToken(t *testing.T) {
t.Parallel()
first := &ecs.RunTaskInput{Count: aws.Int64(2), ClientToken: aws.String("a")}
sameParams := &ecs.RunTaskInput{Count: aws.Int64(2), ClientToken: aws.String("b")}
changed := &ecs.RunTaskInput{Count: aws.Int64(5), ClientToken: aws.String("a")}
assert.Equal(t, runTaskParamHash(first), runTaskParamHash(sameParams),
"the token must not feed the hash")
assert.NotEqual(t, runTaskParamHash(first), runTaskParamHash(changed),
"a changed count must change the hash")
}
// Hashing must not disturb the input, which is dispatched to the launch after.
func TestRunTaskParamHash_LeavesTheInputIntact(t *testing.T) {
t.Parallel()
input := &ecs.RunTaskInput{Count: aws.Int64(2), ClientToken: aws.String("tok")}
runTaskParamHash(input)
require.NotNil(t, input.ClientToken)
assert.Equal(t, "tok", *input.ClientToken)
}
// --- runTaskIdempotent -------------------------------------------------------
func tokenTestInput(token string) *ecs.RunTaskInput {
input := &ecs.RunTaskInput{TaskDefinition: aws.String("web:1")}
if token != "" {
input.ClientToken = aws.String(token)
}
return input
}
func placingService(arn string) *fakeECSService {
out := taskOutput(arn)
return &fakeECSService{describeOut: taskDefWithRole(), runOut: &out}
}
// A request with no token must not pay for the store, let alone fail when the
// store is unreachable: AWS treats an absent token as opting out.
func TestRunTaskIdempotent_NoTokenRunsUnwrapped(t *testing.T) {
t.Parallel()
svc := placingService("task-a")
openStore := func() (*runTaskStore, error) {
t.Error("the store must not be opened for an untokened request")
return nil, nil
}
out, err := runTaskIdempotent(t.Context(), svc, tokenTestAccount, tokenTestInput(""),
allowCheck(t, testTaskRoleARN), openStore)
require.NoError(t, err)
assert.True(t, svc.runCalled, "an untokened request still places")
require.Len(t, out.Tasks, 1)
}
// With no store there is no way to tell a retry from a first attempt, so placing
// anyway would create exactly the duplicate the token was sent to prevent.
func TestRunTaskIdempotent_StoreFailureRefusesToPlace(t *testing.T) {
t.Parallel()
svc := placingService("task-a")
openStore := func() (*runTaskStore, error) { return nil, errors.New("jetstream down") }
_, err := runTaskIdempotent(t.Context(), svc, tokenTestAccount, tokenTestInput("tok-store"),
allowCheck(t, testTaskRoleARN), openStore)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorServerInternal, err.Error())
assert.False(t, svc.runCalled, "a request that cannot be deduplicated must not place")
}
// The whole path, from the token on the input through to a replayed output.
func TestRunTaskIdempotent_PlacesOnceAcrossRetries(t *testing.T) {
t.Parallel()
store := newRunTaskStore(t)
openStore := func() (*runTaskStore, error) { return store, nil }
first := placingService("task-a")
firstOut, err := runTaskIdempotent(t.Context(), first, tokenTestAccount, tokenTestInput("tok-retry"),
allowCheck(t, testTaskRoleARN), openStore)
require.NoError(t, err)
require.Len(t, firstOut.Tasks, 1)
// A retry is a fresh service: the replay must come from the store, not from
// any state the first call left on the fake.
second := placingService("task-b")
replay, err := runTaskIdempotent(t.Context(), second, tokenTestAccount, tokenTestInput("tok-retry"),
allowCheck(t, testTaskRoleARN), openStore)
require.NoError(t, err)
assert.False(t, second.runCalled, "the retry must replay rather than place again")
require.Len(t, replay.Tasks, 1)
assert.Equal(t, "task-a", aws.StringValue(replay.Tasks[0].TaskArn))
}
// A launch error must reach the caller unwrapped, and must not be recorded as a
// placement that a retry would replay.
func TestRunTaskIdempotent_LaunchErrorPropagates(t *testing.T) {
t.Parallel()
store := newRunTaskStore(t)
openStore := func() (*runTaskStore, error) { return store, nil }
svc := &fakeECSService{describeOut: taskDefWithRole()}
_, err := runTaskIdempotent(t.Context(), svc, tokenTestAccount, tokenTestInput("tok-denied"),
denyCheck(t, testTaskRoleARN), openStore)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorAccessDenied, err.Error())
assert.False(t, svc.runCalled)
}
// The store binds once per process, so a second caller must get the first
// caller's store rather than rebind the bucket.
func TestGetRunTaskStore_BindsOnce(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
first, err := getRunTaskStore(t.Context(), nc)
require.NoError(t, err)
require.NotNil(t, first)
second, err := getRunTaskStore(t.Context(), nc)
require.NoError(t, err)
assert.Same(t, first, second, "the bucket must be bound once, not per request")
}
@@ -254,6 +254,9 @@ func (s *LaunchTemplateServiceImpl) resolveVersionNumber(ctx context.Context, ac
// --- CreateLaunchTemplate ---
// ClientToken is accepted but not stored: claimName reserves the name with an
// atomic Create, so a duplicate request cannot make a second template. It gets
// AlreadyExists where AWS would replay the first response.
func (s *LaunchTemplateServiceImpl) CreateLaunchTemplate(ctx context.Context, input *ec2.CreateLaunchTemplateInput, accountID string) (*ec2.CreateLaunchTemplateOutput, error) {
if input.LaunchTemplateData == nil {
return nil, errors.New(awserrors.ErrorMissingParameter)
@@ -433,6 +436,9 @@ func (s *LaunchTemplateServiceImpl) CreateLaunchTemplateVersion(ctx context.Cont
// --- ModifyLaunchTemplate ---
// ClientToken is accepted but not stored: this only moves the default version
// under a CAS, so applying the same request twice converges on the same state
// and creates nothing.
func (s *LaunchTemplateServiceImpl) ModifyLaunchTemplate(ctx context.Context, input *ec2.ModifyLaunchTemplateInput, accountID string) (*ec2.ModifyLaunchTemplateOutput, error) {
header, entry, err := s.resolveHeader(ctx, accountID, input.LaunchTemplateId, input.LaunchTemplateName)
if err != nil {
+1 -1
View File
@@ -123,7 +123,7 @@ func (m *natsEIPManager) Release(ctx context.Context, accountID, allocationID st
// CreateService persists a REPLICA service and reconciles up to desiredCount.
// DAEMON scheduling and serviceRegistries are rejected in v1 (Q15). Re-creating
// an existing ACTIVE service returns the existing record (idempotent).
// an existing ACTIVE service returns it, so ClientToken needs no store.
func (s *Service) CreateService(ctx context.Context, input *ecs.CreateServiceInput, accountID string) (*ecs.CreateServiceOutput, error) {
name := aws.StringValue(input.ServiceName)
if name == "" {