[wanda] Wire context owner override through tar writing (#480)

Add an owner field to tarStream that, when set, overrides the uid/gid on
all file entries written to the tar. Pass the owner through
tarFile.writeTo so each file header gets the override applied.

Topic: add-context-owner-tar
Relative: add-context-owner

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: andrew <andrew@anyscale.com>

---------

Signed-off-by: andrew <andrew@anyscale.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: andrew-anyscale

wanda/tar_file.go

@@ -19,7 +19,9 @@ type tarFile struct {
 	meta *tarMeta
 }
 
-func (t *tarFile) writeTo(tw *tar.Writer, modTime time.Time) error {
+func (t *tarFile) writeTo(
+	tw *tar.Writer, modTime time.Time, owner *contextOwner,
+) error {
 	stat, err := os.Lstat(t.srcFile)
 	if err != nil {
 		return fmt.Errorf("stat file %q: %w", t.srcFile, err)
@@ -37,6 +39,10 @@ func (t *tarFile) writeTo(tw *tar.Writer, modTime time.Time) error {
 		Uid:     meta.UserID,
 		ModTime: modTime,
 	}
+	if owner != nil {
+		hdr.Uid = owner.UserID
+		hdr.Gid = owner.GroupID
+	}
 
 	switch stat.Mode() & os.ModeType {
 	case os.ModeSymlink:
@@ -94,7 +100,7 @@ type tarFileRecord struct {
 	Symlink string `json:"symlink,omitempty"`
 }
 
-func (t *tarFile) record() (*tarFileRecord, error) {
+func (t *tarFile) record(owner *contextOwner) (*tarFileRecord, error) {
 	// Use Lstat to detect symlinks without following them.
 	stat, err := os.Lstat(t.srcFile)
 	if err != nil {
@@ -125,15 +131,20 @@ func (t *tarFile) record() (*tarFileRecord, error) {
 		h.Write([]byte(target))
 		contentDigest := sha256DigestString(h)
 
-		return &tarFileRecord{
+		r := &tarFileRecord{
 			Name:          t.name,
 			Mode:          mode,
 			GroupID:       meta.GroupID,
 			UserID:        meta.UserID,
 			Size:          0, // Symlinks have no content size in tar.
 			ContentDigest: contentDigest,
 			Symlink:       target,
-		}, nil
+		}
+		if owner != nil {
+			r.UserID = owner.UserID
+			r.GroupID = owner.GroupID
+		}
+		return r, nil
 	default:
 		f, err := os.Open(t.srcFile)
 		if err != nil {
@@ -156,7 +167,10 @@ func (t *tarFile) record() (*tarFileRecord, error) {
 
 			ContentDigest: contentDigest,
 		}
-
+		if owner != nil {
+			r.UserID = owner.UserID
+			r.GroupID = owner.GroupID
+		}
 		return r, nil
 	}
 }

wanda/tar_file_test.go

@@ -38,7 +38,7 @@ func TestTarFile(t *testing.T) {
 	// Tarball only stores the mod time with second precision.
 	modTime := time.Now().Truncate(time.Second)
 
-	if err := tf.writeTo(tw, modTime); err != nil {
+	if err := tf.writeTo(tw, modTime, nil); err != nil {
 		t.Fatalf("write to tar stream: %v", err)
 	}
 	if err := tw.Close(); err != nil {
@@ -100,7 +100,7 @@ func TestTarFileRecord(t *testing.T) {
 		},
 	}
 
-	r, err := tf.record()
+	r, err := tf.record(nil)
 	if err != nil {
 		t.Fatalf("record: %v", err)
 	}
@@ -142,7 +142,7 @@ func TestTarFileRecord(t *testing.T) {
 		meta:    tf.meta, // The same meta as the first file.
 	}
 
-	r2, err := tf2.record()
+	r2, err := tf2.record(nil)
 	if err != nil {
 		t.Fatalf("record for file 2: %v", err)
 	}
@@ -185,7 +185,7 @@ func TestTarFileSymlink(t *testing.T) {
 	tw := tar.NewWriter(tarBuf)
 	modTime := time.Now().Truncate(time.Second)
 
-	if err := tf.writeTo(tw, modTime); err != nil {
+	if err := tf.writeTo(tw, modTime, nil); err != nil {
 		t.Fatalf("write symlink to tar stream: %v", err)
 	}
 	if err := tw.Close(); err != nil {
@@ -234,7 +234,7 @@ func TestTarFileSymlinkRecord(t *testing.T) {
 		},
 	}
 
-	r, err := tf.record()
+	r, err := tf.record(nil)
 	if err != nil {
 		t.Fatalf("record: %v", err)
 	}
@@ -277,7 +277,7 @@ func TestTarFileSymlinkRecord(t *testing.T) {
 		meta:    tf.meta,
 	}
 
-	r2, err := tf2.record()
+	r2, err := tf2.record(nil)
 	if err != nil {
 		t.Fatalf("record for link2: %v", err)
 	}

wanda/tar_stream.go

@@ -18,6 +18,9 @@ type tarStream struct {
 	// all files will use this mod time default. This makes the stream
 	// deterministic and cachable.
 	modTime time.Time
+
+	// owner overrides the uid/gid for all entries when set.
+	owner *contextOwner
 }
 
 // newTarStream creates a new tarball stream.
@@ -57,8 +60,10 @@ func (s *tarStream) writeTo(tw *tar.Writer) error {
 
 	for _, name := range names {
 		f := s.files[name]
-		if err := f.writeTo(tw, s.modTime); err != nil {
-			return fmt.Errorf("write file %q to stream: %w", name, err)
+		if err := f.writeTo(tw, s.modTime, s.owner); err != nil {
+			return fmt.Errorf(
+				"write file %q to stream: %w", name, err,
+			)
 		}
 	}
 	return nil
@@ -90,7 +95,7 @@ func (s *tarStream) digest() (string, error) {
 	for _, name := range names {
 		f := s.files[name]
 
-		r, err := f.record()
+		r, err := f.record(s.owner)
 		if err != nil {
 			return "", fmt.Errorf("digest file %q: %w", name, err)
 		}

wanda/tar_stream_test.go

@@ -9,6 +9,54 @@ import (
 	"testing"
 )
 
+func TestTarStreamContextOwner(t *testing.T) {
+	tmp := t.TempDir()
+
+	f := filepath.Join(tmp, "testfile")
+	if err := os.WriteFile(f, []byte("data"), 0644); err != nil {
+		t.Fatalf("write file: %v", err)
+	}
+
+	ts := newTarStream()
+	ts.owner = &contextOwner{UserID: 2000, GroupID: 100}
+	ts.addFile("a/b/file.txt", nil, f)
+
+	r := newWriterToReader(ts)
+	buf := new(bytes.Buffer)
+	if _, err := io.Copy(buf, r); err != nil {
+		t.Fatalf("copy tar stream: %v", err)
+	}
+
+	tr := tar.NewReader(buf)
+
+	hdr, err := tr.Next()
+	if err != nil {
+		t.Fatalf("read file header: %v", err)
+	}
+	if hdr.Name != "a/b/file.txt" {
+		t.Errorf("got name %q, want %q", hdr.Name, "a/b/file.txt")
+	}
+	if hdr.Uid != 2000 {
+		t.Errorf("got uid %d, want %d", hdr.Uid, 2000)
+	}
+	if hdr.Gid != 100 {
+		t.Errorf("got gid %d, want %d", hdr.Gid, 100)
+	}
+
+	d1, err := ts.digest()
+	if err != nil {
+		t.Fatalf("digest: %v", err)
+	}
+	ts.owner = &contextOwner{UserID: 3000, GroupID: 300}
+	d2, err := ts.digest()
+	if err != nil {
+		t.Fatalf("digest: %v", err)
+	}
+	if d1 == d2 {
+		t.Error("digest should change when owner is modified")
+	}
+}
+
 func TestTarStream(t *testing.T) {
 	tmp := t.TempDir()