snapshots: stamp the pod's caller annotations on hibernate and latest pushes

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.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

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.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=

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,11 +47,20 @@ 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)
 	}
 	metrics.SnapshotPushDuration.WithLabelValues(pod.Namespace).Observe(time.Since(pushStart).Seconds())
 	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
+}

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)

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,

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()
+}