3 files changed, 80 insertions, 13 deletions

diff --git a/pkg/btrfs/io1_device.go b/pkg/btrfs/io1_device.go
index 3b84fcc..978d388 100644
--- a/pkg/btrfs/io1_device.go
+++ b/pkg/btrfs/io1_device.go
@@ -33,6 +33,10 @@ func (dev *Device) ReadAt(dat []byte, paddr PhysicalAddr) (int, error) {
 	return dev.File.ReadAt(dat, int64(paddr))
 }
 
+func (dev *Device) WriteAt(dat []byte, paddr PhysicalAddr) (int, error) {
+	return dev.File.WriteAt(dat, int64(paddr))
+}
+
 func (dev *Device) Superblocks() ([]*util.Ref[PhysicalAddr, Superblock], error) {
 	if dev.cacheSuperblocks != nil {
 		return dev.cacheSuperblocks, nil
diff --git a/pkg/btrfs/io2_fs.go b/pkg/btrfs/io2_fs.go
index 035019c..5d1e343 100644
--- a/pkg/btrfs/io2_fs.go
+++ b/pkg/btrfs/io2_fs.go
@@ -153,18 +153,6 @@ func (fs *FS) Init() error {
 	return nil
 }
 
-func (fs *FS) ReadAt(dat []byte, laddr LogicalAddr) (int, error) {
-	done := 0
-	for done < len(dat) {
-		n, err := fs.maybeShortReadAt(dat[done:], laddr+LogicalAddr(done))
-		done += n
-		if err != nil {
-			return done, err
-		}
-	}
-	return done, nil
-}
-
 type QualifiedPhysicalAddr struct {
 	Dev  UUID
 	Addr PhysicalAddr
@@ -190,6 +178,18 @@ func (fs *FS) Resolve(laddr LogicalAddr) (paddrs map[QualifiedPhysicalAddr]struc
 	return paddrs, maxlen
 }
 
+func (fs *FS) ReadAt(dat []byte, laddr LogicalAddr) (int, error) {
+	done := 0
+	for done < len(dat) {
+		n, err := fs.maybeShortReadAt(dat[done:], laddr+LogicalAddr(done))
+		done += n
+		if err != nil {
+			return done, err
+		}
+	}
+	return done, nil
+}
+
 func (fs *FS) maybeShortReadAt(dat []byte, laddr LogicalAddr) (int, error) {
 	paddrs, maxlen := fs.Resolve(laddr)
 	if len(paddrs) == 0 {
@@ -219,3 +219,36 @@ func (fs *FS) maybeShortReadAt(dat []byte, laddr LogicalAddr) (int, error) {
 	}
 	return len(dat), nil
 }
+
+func (fs *FS) WriteAt(dat []byte, laddr LogicalAddr) (int, error) {
+	done := 0
+	for done < len(dat) {
+		n, err := fs.maybeShortWriteAt(dat[done:], laddr+LogicalAddr(done))
+		done += n
+		if err != nil {
+			return done, err
+		}
+	}
+	return done, nil
+}
+
+func (fs *FS) maybeShortWriteAt(dat []byte, laddr LogicalAddr) (int, error) {
+	paddrs, maxlen := fs.Resolve(laddr)
+	if len(paddrs) == 0 {
+		return 0, fmt.Errorf("could not map logical address %v", laddr)
+	}
+	if uint64(len(dat)) > maxlen {
+		dat = dat[:maxlen]
+	}
+
+	for paddr := range paddrs {
+		dev, ok := fs.uuid2dev[paddr.Dev]
+		if !ok {
+			return 0, fmt.Errorf("device=%s does not exist", paddr.Dev)
+		}
+		if _, err := dev.WriteAt(dat, paddr.Addr); err != nil {
+			return 0, fmt.Errorf("write device=%s paddr=%v: %w", paddr.Dev, paddr.Addr, err)
+		}
+	}
+	return len(dat), nil
+}
diff --git a/pkg/util/ref.go b/pkg/util/ref.go
index 69f7db4..05ef08f 100644
--- a/pkg/util/ref.go
+++ b/pkg/util/ref.go
@@ -1,6 +1,9 @@
 package util
 
 import (
+	"fmt"
+	"io"
+
 	"lukeshu.com/btrfs-tools/pkg/binstruct"
 )
 
@@ -8,8 +11,14 @@ type File[A ~int64] interface {
 	Name() string
 	Size() (A, error)
 	ReadAt(p []byte, off A) (n int, err error)
+	WriteAt(p []byte, off A) (n int, err error)
 }
 
+var (
+	_ io.WriterAt = File[int64](nil)
+	_ io.ReaderAt = File[int64](nil)
+)
+
 type Ref[A ~int64, T any] struct {
 	File File[A]
 	Addr A
@@ -22,7 +31,28 @@ func (r *Ref[A, T]) Read() error {
 	if _, err := r.File.ReadAt(buf, r.Addr); err != nil {
 		return err
 	}
-	if _, err := binstruct.Unmarshal(buf, &r.Data); err != nil {
+	n, err := binstruct.Unmarshal(buf, &r.Data)
+	if err != nil {
+		return err
+	}
+	if n != size {
+		return fmt.Errorf("util.Ref[%T].Read: left over data: read %d bytes but only consumed %d",
+			r.Data, size, n)
+	}
+	return nil
+}
+
+func (r *Ref[A, T]) Write() error {
+	size := binstruct.StaticSize(r.Data)
+	buf, err := binstruct.Marshal(r.Data)
+	if err != nil {
+		return err
+	}
+	if len(buf) != size {
+		return fmt.Errorf("util.Ref[%T].Write: expected to want to write %d bytes, but got %d bytes to write",
+			r.Data, size, len(buf))
+	}
+	if _, err = r.File.WriteAt(buf, r.Addr); err != nil {
 		return err
 	}
 	return nil