From 45f3581b02eff4dc4a0747e6b8ad4339caad81dc Mon Sep 17 00:00:00 2001 From: CMGS Date: Tue, 29 Sep 2026 00:46:12 +0800 Subject: [PATCH 1/2] snapshots: stamp the pod's caller annotations on hibernate and latest pushes Reads vm.cocoonstack.io/snapshot-annotations through cocoon-common's typed reader and passes the map to the pusher, which keeps the peer-discovery node hint last so a caller cannot override it. A malformed value is logged and the push goes ahead without caller annotations: metadata never blocks a hibernate or a delete. Pins cocoon-common to the branch commit of cocoon-common#36. --- go.mod | 2 +- go.sum | 4 +- provider/cocoon/snapshot.go | 12 +++++- provider/cocoon/update.go | 2 +- snapshots/pusher.go | 12 +++--- snapshots/pusher_test.go | 84 ++++++++++++++++++++++++++++++++++++- 6 files changed, 105 insertions(+), 11 deletions(-) diff --git a/go.mod b/go.mod index 1ebdea3..343a6f4 100644 --- a/go.mod +++ b/go.mod @@ -3,7 +3,7 @@ module github.com/cocoonstack/vk-cocoon go 1.27.1 require ( - github.com/cocoonstack/cocoon-common v0.3.6 + github.com/cocoonstack/cocoon-common v0.3.7-0.20260928164220-c33e27e571b8 github.com/google/go-containerregistry v0.22.1 github.com/projecteru2/core v0.1.5 github.com/prometheus/client_golang v1.24.1 diff --git a/go.sum b/go.sum index e22bb8a..cdd6a85 100644 --- a/go.sum +++ b/go.sum @@ -30,8 +30,8 @@ github.com/cockroachdb/logtags v0.0.0-20241215232642-bb51bb14a506 h1:ASDL+UJcILM github.com/cockroachdb/logtags v0.0.0-20241215232642-bb51bb14a506/go.mod h1:Mw7HqKr2kdtu6aYGn3tPmAftiP3QPX63LdK/zcariIo= github.com/cockroachdb/redact v1.1.8 h1:8eVLLj6juKxiKrAEw2b8cJvNqWq++U8WOfQFuL7KTaA= github.com/cockroachdb/redact v1.1.8/go.mod h1:GceHHpJ0rMDpYARL5In88Alq/xMBUtVlz7Qxix6ZVkw= -github.com/cocoonstack/cocoon-common v0.3.6 h1:iRrIyeRXHgYuwKYDOB6TjidiZL1oeyfOYkTGXDSD/PU= -github.com/cocoonstack/cocoon-common v0.3.6/go.mod h1:Gw/zBYqntoxpcT//2rBcXHQz6JErin/IsTsEmoQdGr0= +github.com/cocoonstack/cocoon-common v0.3.7-0.20260928164220-c33e27e571b8 h1:BsweRkm/ERnUHeSnFhgfGd9tCfU2H6igjqAe7BRxrns= +github.com/cocoonstack/cocoon-common v0.3.7-0.20260928164220-c33e27e571b8/go.mod h1:Gw/zBYqntoxpcT//2rBcXHQz6JErin/IsTsEmoQdGr0= github.com/coreos/go-semver v0.3.1 h1:yi21YpKnrx1gt5R+la8n5WgS0kCrsPp33dmEyHReZr4= github.com/coreos/go-semver v0.3.1/go.mod h1:irMmmIw/7yzSRPWryHsK7EYSg09caPQL03VsM8rvUec= github.com/coreos/go-systemd/v22 v22.7.0 h1:LAEzFkke61DFROc7zNLX/WA2i5J8gYqe0rSj9KI28KA= diff --git a/provider/cocoon/snapshot.go b/provider/cocoon/snapshot.go index 747104a..dbbc43d 100644 --- a/provider/cocoon/snapshot.go +++ b/provider/cocoon/snapshot.go @@ -9,6 +9,7 @@ import ( "github.com/projecteru2/core/log" corev1 "k8s.io/api/core/v1" + "github.com/cocoonstack/cocoon-common/meta" "github.com/cocoonstack/vk-cocoon/metrics" "github.com/cocoonstack/vk-cocoon/vm" ) @@ -46,7 +47,7 @@ func (p *Provider) saveAndPushSnapshot(ctx context.Context, pod *corev1.Pod, v * metrics.SnapshotSaveTotal.WithLabelValues("ok").Inc() pushStart := time.Now() - if err := p.Pusher.PushSnapshot(ctx, v.Name, "", tag, image); err != nil { + if err := p.Pusher.PushSnapshot(ctx, v.Name, "", tag, image, snapshotAnnotations(ctx, pod)); err != nil { metrics.SnapshotPushTotal.WithLabelValues("failed").Inc() return fmt.Errorf("push snapshot %s: %w", v.Name, err) } @@ -54,3 +55,12 @@ func (p *Provider) saveAndPushSnapshot(ctx context.Context, pod *corev1.Pod, v * metrics.SnapshotPushTotal.WithLabelValues("ok").Inc() return nil } + +// snapshotAnnotations returns the pod's caller annotations for a push; a malformed value pushes without them. +func snapshotAnnotations(ctx context.Context, pod *corev1.Pod) map[string]string { + annotations, err := meta.ReadSnapshotAnnotations(pod) + if err != nil { + log.WithFunc("cocoon.snapshotAnnotations").Warnf(ctx, "push snapshot for %s/%s without caller annotations: %v", pod.Namespace, pod.Name, err) + } + return annotations +} diff --git a/provider/cocoon/update.go b/provider/cocoon/update.go index c601aec..ffef988 100644 --- a/provider/cocoon/update.go +++ b/provider/cocoon/update.go @@ -188,7 +188,7 @@ func (p *Provider) hibernate(ctx context.Context, pod *corev1.Pod, spec meta.VMS if p.Pusher != nil { pushStart := time.Now() // spec.Image rides as the baseimage annotation for hibernateEvidence's identity guard. - if err := p.Pusher.PushSnapshot(ctx, v.Name, v.Name, meta.HibernateSnapshotTag, spec.Image); err != nil { + if err := p.Pusher.PushSnapshot(ctx, v.Name, v.Name, meta.HibernateSnapshotTag, spec.Image, snapshotAnnotations(ctx, pod)); err != nil { metrics.SnapshotPushTotal.WithLabelValues("failed").Inc() metrics.HibernateTotal.WithLabelValues(pod.Namespace, "push", "failed").Inc() p.rollbackHibernate(ctx, pod, v, dropNIC) diff --git a/snapshots/pusher.go b/snapshots/pusher.go index 6e34009..3ed0226 100644 --- a/snapshots/pusher.go +++ b/snapshots/pusher.go @@ -4,6 +4,7 @@ import ( "cmp" "context" "fmt" + "maps" "golang.org/x/sync/semaphore" @@ -27,8 +28,8 @@ type Pusher struct { NodeName string } -// PushSnapshot uploads a snapshot to the registry at the given repo/tag. -func (p *Pusher) PushSnapshot(ctx context.Context, vmName, repo, tag, baseImage string) error { +// PushSnapshot uploads a snapshot to the registry at the given repo/tag with the caller's annotations. +func (p *Pusher) PushSnapshot(ctx context.Context, vmName, repo, tag, baseImage string, annotations map[string]string) error { repo = cmp.Or(repo, vmName) tag = cmp.Or(tag, meta.DefaultSnapshotTag) @@ -42,15 +43,16 @@ func (p *Pusher) PushSnapshot(ctx context.Context, vmName, repo, tag, baseImage Cocoon: runnerAdapter{Runtime: p.Runtime}, } - var annotations map[string]string + stamped := make(map[string]string, len(annotations)+1) + maps.Copy(stamped, annotations) if p.NodeName != "" { - annotations = map[string]string{AnnotationFromNode: p.NodeName} + stamped[AnnotationFromNode] = p.NodeName } if err := pusher.Push(ctx, snapshot.PushOptions{ Name: repo, Tag: tag, BaseImage: baseImage, - Annotations: annotations, + Annotations: stamped, ZstdLevel: p.Transfer.ZstdLevel, ChunkSizeMiB: p.Transfer.ChunkSizeMiB, Concurrency: p.Transfer.Concurrency, diff --git a/snapshots/pusher_test.go b/snapshots/pusher_test.go index 2e45800..ec2ae6d 100644 --- a/snapshots/pusher_test.go +++ b/snapshots/pusher_test.go @@ -1,10 +1,19 @@ package snapshots import ( + "archive/tar" + "bytes" "context" + "encoding/json" "errors" + "io" "testing" "time" + + "github.com/cocoonstack/cocoon-common/manifest" + "github.com/cocoonstack/cocoon-common/oci" + "github.com/cocoonstack/cocoon-common/snapshot" + "github.com/cocoonstack/vk-cocoon/vm" ) func TestPushSnapshotWaitsOnTheGateForSpoolPushes(t *testing.T) { @@ -14,8 +23,81 @@ func TestPushSnapshotWaitsOnTheGateForSpoolPushes(t *testing.T) { defer pushGate.Release(1) ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) defer cancel() - err := (&Pusher{}).PushSnapshot(ctx, "vm-a", "", "", "") + err := (&Pusher{}).PushSnapshot(ctx, "vm-a", "", "", "", nil) if !errors.Is(err, context.DeadlineExceeded) { t.Fatalf("a v1 spool push with the gate held returned %v, want the gate wait to expire", err) } } + +func TestPushSnapshotStampsCallerAnnotationsBesideTheNodeHint(t *testing.T) { + reg := &manifestRecorder{} + p := &Pusher{Registry: reg, Runtime: exportRuntime{export: exportTar(t, "vm-a")}, NodeName: "node-b"} + caller := map[string]string{"ai.simular.owner": "org:osworld", AnnotationFromNode: "spoofed"} + if err := p.PushSnapshot(t.Context(), "vm-a", "", "hibernate", "img:1", caller); err != nil { + t.Fatalf("push: %v", err) + } + var m struct { + Annotations map[string]string `json:"annotations"` + } + if err := json.Unmarshal(reg.manifest, &m); err != nil { + t.Fatalf("decode manifest: %v", err) + } + if got := m.Annotations["ai.simular.owner"]; got != "org:osworld" { + t.Errorf("caller annotation = %q, want org:osworld", got) + } + if got := m.Annotations[AnnotationFromNode]; got != "node-b" { + t.Errorf("%s = %q, want the pushing node", AnnotationFromNode, got) + } + if got := m.Annotations[manifest.AnnotationSnapshotBaseImage]; got != "img:1" { + t.Errorf("%s = %q, want img:1", manifest.AnnotationSnapshotBaseImage, got) + } + if caller[AnnotationFromNode] != "spoofed" { + t.Errorf("push mutated the caller's map") + } +} + +type manifestRecorder struct { + oci.Registry + manifest []byte +} + +func (r *manifestRecorder) HasBlob(context.Context, string, string) (bool, error) { return false, nil } + +func (r *manifestRecorder) PutBlob(_ context.Context, _, _ string, body io.ReadSeeker, _ int64) error { + _, err := io.Copy(io.Discard, body) + return err +} + +func (r *manifestRecorder) PutManifest(_ context.Context, _, _ string, data []byte, _ string) error { + r.manifest = data + return nil +} + +type exportRuntime struct { + vm.Runtime + export []byte +} + +func (r exportRuntime) SnapshotExport(context.Context, string) (io.ReadCloser, func() error, error) { + return io.NopCloser(bytes.NewReader(r.export)), func() error { return nil }, nil +} + +func exportTar(t *testing.T, name string) []byte { + t.Helper() + envelope, err := snapshot.MarshalEnvelope(&manifest.SnapshotConfig{SchemaVersion: "v1", SnapshotID: "SNAP-1"}, name) + if err != nil { + t.Fatal(err) + } + var buf bytes.Buffer + tw := tar.NewWriter(&buf) + if err := tw.WriteHeader(&tar.Header{Name: "snapshot.json", Mode: 0o644, Size: int64(len(envelope))}); err != nil { + t.Fatal(err) + } + if _, err := tw.Write(envelope); err != nil { + t.Fatal(err) + } + if err := tw.Close(); err != nil { + t.Fatal(err) + } + return buf.Bytes() +} From d7388e055c4b27dc35e6120cc0bd7143ca07f3a5 Mon Sep 17 00:00:00 2001 From: CMGS Date: Tue, 29 Sep 2026 01:28:44 +0800 Subject: [PATCH 2/2] build: pin cocoon-common to main b03e52d --- go.mod | 2 +- go.sum | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/go.mod b/go.mod index 343a6f4..a038df0 100644 --- a/go.mod +++ b/go.mod @@ -3,7 +3,7 @@ module github.com/cocoonstack/vk-cocoon go 1.27.1 require ( - github.com/cocoonstack/cocoon-common v0.3.7-0.20260928164220-c33e27e571b8 + github.com/cocoonstack/cocoon-common v0.3.7-0.20260928172801-b03e52d1f721 github.com/google/go-containerregistry v0.22.1 github.com/projecteru2/core v0.1.5 github.com/prometheus/client_golang v1.24.1 diff --git a/go.sum b/go.sum index cdd6a85..dade226 100644 --- a/go.sum +++ b/go.sum @@ -30,8 +30,8 @@ github.com/cockroachdb/logtags v0.0.0-20241215232642-bb51bb14a506 h1:ASDL+UJcILM github.com/cockroachdb/logtags v0.0.0-20241215232642-bb51bb14a506/go.mod h1:Mw7HqKr2kdtu6aYGn3tPmAftiP3QPX63LdK/zcariIo= github.com/cockroachdb/redact v1.1.8 h1:8eVLLj6juKxiKrAEw2b8cJvNqWq++U8WOfQFuL7KTaA= github.com/cockroachdb/redact v1.1.8/go.mod h1:GceHHpJ0rMDpYARL5In88Alq/xMBUtVlz7Qxix6ZVkw= -github.com/cocoonstack/cocoon-common v0.3.7-0.20260928164220-c33e27e571b8 h1:BsweRkm/ERnUHeSnFhgfGd9tCfU2H6igjqAe7BRxrns= -github.com/cocoonstack/cocoon-common v0.3.7-0.20260928164220-c33e27e571b8/go.mod h1:Gw/zBYqntoxpcT//2rBcXHQz6JErin/IsTsEmoQdGr0= +github.com/cocoonstack/cocoon-common v0.3.7-0.20260928172801-b03e52d1f721 h1:xUtBRurq3bzpL2D0P0Apeiq4jjplNLOD/Kbd6G14WwM= +github.com/cocoonstack/cocoon-common v0.3.7-0.20260928172801-b03e52d1f721/go.mod h1:Gw/zBYqntoxpcT//2rBcXHQz6JErin/IsTsEmoQdGr0= github.com/coreos/go-semver v0.3.1 h1:yi21YpKnrx1gt5R+la8n5WgS0kCrsPp33dmEyHReZr4= github.com/coreos/go-semver v0.3.1/go.mod h1:irMmmIw/7yzSRPWryHsK7EYSg09caPQL03VsM8rvUec= github.com/coreos/go-systemd/v22 v22.7.0 h1:LAEzFkke61DFROc7zNLX/WA2i5J8gYqe0rSj9KI28KA=