diff --git a/cmd/atelet/local_checkpoints.go b/cmd/atelet/local_checkpoints.go new file mode 100644 index 000000000..5426cbc2b --- /dev/null +++ b/cmd/atelet/local_checkpoints.go @@ -0,0 +1,54 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "context" + "log/slog" + "os" + "path/filepath" + + "github.com/agent-substrate/substrate/internal/ateompath" +) + +// pruneLocalCheckpoints removes the actor's local pause snapshots except keep; +// empty keep removes all. Best-effort: failures are logged, never fatal. +func pruneLocalCheckpoints(ctx context.Context, actorUID, keep string) { + pruneLocalCheckpointDir(ctx, ateompath.LocalCheckpointsDir(actorUID), keep) +} + +func pruneLocalCheckpointDir(ctx context.Context, dir, keep string) { + entries, err := os.ReadDir(dir) + if err != nil { + if !os.IsNotExist(err) { + slog.WarnContext(ctx, "failed to list local checkpoints for pruning", slog.String("dir", dir), slog.Any("err", err)) + } + return + } + for _, entry := range entries { + if keep != "" && entry.Name() == keep { + continue + } + path := filepath.Join(dir, entry.Name()) + if err := os.RemoveAll(path); err != nil { + slog.WarnContext(ctx, "failed to prune local checkpoint", slog.String("path", path), slog.Any("err", err)) + continue + } + slog.InfoContext(ctx, "pruned local checkpoint", slog.String("path", path)) + } + if keep == "" { + _ = os.Remove(dir) + } +} diff --git a/cmd/atelet/local_checkpoints_test.go b/cmd/atelet/local_checkpoints_test.go new file mode 100644 index 000000000..c7503559f --- /dev/null +++ b/cmd/atelet/local_checkpoints_test.go @@ -0,0 +1,78 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "context" + "os" + "path/filepath" + "testing" +) + +func writeSnapshotDir(t *testing.T, dir, prefix string) { + t.Helper() + p := filepath.Join(dir, prefix) + if err := os.MkdirAll(p, 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(p, "memory.img"), []byte("x"), 0o600); err != nil { + t.Fatal(err) + } +} + +func listDir(t *testing.T, dir string) []string { + t.Helper() + entries, err := os.ReadDir(dir) + if os.IsNotExist(err) { + return nil + } + if err != nil { + t.Fatal(err) + } + var names []string + for _, e := range entries { + names = append(names, e.Name()) + } + return names +} + +func TestPruneKeepsOnlyNamedSnapshot(t *testing.T) { + dir := t.TempDir() + writeSnapshotDir(t, dir, "pause-1") + writeSnapshotDir(t, dir, "pause-2") + writeSnapshotDir(t, dir, "pause-3") + + pruneLocalCheckpointDir(context.Background(), dir, "pause-3") + + if got := listDir(t, dir); len(got) != 1 || got[0] != "pause-3" { + t.Fatalf("dir = %v, want only pause-3", got) + } +} + +func TestPruneEmptyKeepRemovesEverything(t *testing.T) { + dir := t.TempDir() + writeSnapshotDir(t, dir, "pause-1") + writeSnapshotDir(t, dir, "pause-2") + + pruneLocalCheckpointDir(context.Background(), dir, "") + + if _, err := os.Stat(dir); !os.IsNotExist(err) { + t.Fatalf("dir still exists (err=%v), want removed entirely", err) + } +} + +func TestPruneMissingDirIsNoop(t *testing.T) { + pruneLocalCheckpointDir(context.Background(), filepath.Join(t.TempDir(), "absent"), "anything") +} diff --git a/cmd/atelet/main.go b/cmd/atelet/main.go index 7b5464699..61e7a9e77 100644 --- a/cmd/atelet/main.go +++ b/cmd/atelet/main.go @@ -398,10 +398,14 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe if err := s.uploadExternalCheckpoint(ctx, req, checkpointDir, sandboxRec); err != nil { return nil, ateerrors.NewGRPCError(ctx, codes.DataLoss, ateerrors.ReasonFaileSaveSnapshot, ateerrors.ActorCrashedMetadata(), fmt.Errorf("%w: while uploading external snapshot: %w", ateerrors.ReasonFaileSaveSnapshot, err)) } + // The durable snapshot supersedes local pause snapshots. + pruneLocalCheckpoints(ctx, actorUID, "") case ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL: if err := s.moveLocalCheckpoint(ctx, req, checkpointDir, sandboxRec); err != nil { return nil, ateerrors.NewGRPCError(ctx, codes.DataLoss, ateerrors.ReasonFaileSaveSnapshot, ateerrors.ActorCrashedMetadata(), fmt.Errorf("%w: while moving to local snapshot: %w", ateerrors.ReasonFaileSaveSnapshot, err)) } + // Without pruning, every pause leaves a snapshot behind until the disk fills. + pruneLocalCheckpoints(ctx, actorUID, req.GetLocalConfig().GetSnapshotPrefix()) default: return nil, fmt.Errorf("unexpected checkpoint type: %v", req.GetType()) } @@ -1049,8 +1053,8 @@ func validateCheckpointRequest(req *ateletpb.CheckpointRequest) error { return err } case ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL: - if req.GetLocalConfig().GetSnapshotPrefix() == "" { - return fmt.Errorf("snapshot prefix must be non-empty for type %s", req.GetType().String()) + if err := resources.ValidateLocalSnapshotPrefix(req.GetLocalConfig().GetSnapshotPrefix()); err != nil { + return err } default: return fmt.Errorf("invalid checkpoint type: %v", req.GetType()) @@ -1101,8 +1105,8 @@ func validateRestoreRequest(req *ateletpb.RestoreRequest) error { return err } case ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL: - if req.GetLocalConfig().GetSnapshotPrefix() == "" { - return fmt.Errorf("snapshot prefix must be non-empty for type %s", req.GetType().String()) + if err := resources.ValidateLocalSnapshotPrefix(req.GetLocalConfig().GetSnapshotPrefix()); err != nil { + return err } default: return fmt.Errorf("invalid checkpoint type: %v", req.GetType()) diff --git a/cmd/atelet/main_test.go b/cmd/atelet/main_test.go index 0378c71c2..fedfb9ef1 100644 --- a/cmd/atelet/main_test.go +++ b/cmd/atelet/main_test.go @@ -277,6 +277,14 @@ func TestValidateCheckpointRequest(t *testing.T) { r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL r.Config = &ateletpb.CheckpointRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: ""}} }), true}, + {"nested local snapshot prefix", makeReq(func(r *ateletpb.CheckpointRequest) { + r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL + r.Config = &ateletpb.CheckpointRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: "pause/2"}} + }), true}, + {"traversal local snapshot prefix", makeReq(func(r *ateletpb.CheckpointRequest) { + r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL + r.Config = &ateletpb.CheckpointRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: ".."}} + }), true}, {"unspecified snapshot type", makeReq(func(r *ateletpb.CheckpointRequest) { r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_UNSPECIFIED }), true}, {"unspecified snapshot scope", makeReq(func(r *ateletpb.CheckpointRequest) { r.Scope = ateletpb.SnapshotScope_SNAPSHOT_SCOPE_UNSPECIFIED }), true}, {"invalid snapshot scope", makeReq(func(r *ateletpb.CheckpointRequest) { r.Scope = ateletpb.SnapshotScope(23) }), true}, @@ -323,6 +331,14 @@ func TestValidateRestoreRequest(t *testing.T) { r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL r.Config = &ateletpb.RestoreRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: ""}} }), true}, + {"nested local snapshot prefix", makeReq(func(r *ateletpb.RestoreRequest) { + r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL + r.Config = &ateletpb.RestoreRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: "pause/2"}} + }), true}, + {"traversal local snapshot prefix", makeReq(func(r *ateletpb.RestoreRequest) { + r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL + r.Config = &ateletpb.RestoreRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: ".."}} + }), true}, {"unspecified snapshot type", makeReq(func(r *ateletpb.RestoreRequest) { r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_UNSPECIFIED }), true}, {"unspecified snapshot scope", makeReq(func(r *ateletpb.RestoreRequest) { r.Scope = ateletpb.SnapshotScope_SNAPSHOT_SCOPE_UNSPECIFIED }), true}, {"invalid snapshot scope", makeReq(func(r *ateletpb.RestoreRequest) { r.Scope = ateletpb.SnapshotScope(23) }), true}, diff --git a/internal/resources/validate.go b/internal/resources/validate.go index 09c985df7..302f8dedd 100644 --- a/internal/resources/validate.go +++ b/internal/resources/validate.go @@ -173,6 +173,20 @@ func ValidateSnapshotURIPrefix(prefix string) error { return nil } +// ValidateLocalSnapshotPrefix ensures a local snapshot prefix is a single +// path segment: it is joined onto the actor's checkpoint directory and +// compared against that directory's entries when pruning, so a nested or +// relative prefix would escape the directory or delete the wrong snapshots. +func ValidateLocalSnapshotPrefix(prefix string) error { + if prefix == "" { + return fmt.Errorf("snapshot prefix must be non-empty") + } + if prefix == "." || prefix == ".." || strings.ContainsAny(prefix, `/\`) { + return fmt.Errorf("invalid snapshot prefix %q: must be a single path segment", prefix) + } + return nil +} + // ValidateWorker checks that the worker message is well-formed. func ValidateWorker(worker *ateapipb.Worker, fldPath *field.Path) field.ErrorList { var errs field.ErrorList diff --git a/internal/resources/validate_test.go b/internal/resources/validate_test.go index b25de1d55..f8a7fad32 100644 --- a/internal/resources/validate_test.go +++ b/internal/resources/validate_test.go @@ -238,6 +238,31 @@ func TestValidateSnapshotURIPrefix(t *testing.T) { } } +func TestValidateLocalSnapshotPrefix(t *testing.T) { + tests := []struct { + name string + prefix string + wantErr bool + }{ + {"valid", "pause", false}, + {"valid with dash and digits", "pause-2", false}, + {"empty", "", true}, + {"dot", ".", true}, + {"dotdot", "..", true}, + {"nested", "pause/2", true}, + {"absolute", "/pause", true}, + {"traversal", "../other-actor", true}, + {"backslash", `pause\2`, true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if err := ValidateLocalSnapshotPrefix(tt.prefix); (err != nil) != tt.wantErr { + t.Errorf("ValidateLocalSnapshotPrefix(%q) err = %v, wantErr %v", tt.prefix, err, tt.wantErr) + } + }) + } +} + func TestValidateWorker(t *testing.T) { tests := []struct { name string