Add an AddrDelta type for safer arithmentic

6 files changed, 33 insertions, 17 deletions

diff --git a/pkg/btrfs/btrfsitem/item_blockgroup.go b/pkg/btrfs/btrfsitem/item_blockgroup.go
index a28d4bf..737cfa0 100644
--- a/pkg/btrfs/btrfsitem/item_blockgroup.go
+++ b/pkg/btrfs/btrfsitem/item_blockgroup.go
@@ -6,6 +6,8 @@ import (
 	"lukeshu.com/btrfs-tools/pkg/util"
 )
 
+// key.objectid = logical_addr
+// key.offset = size of chunk
 type BlockGroup struct { // BLOCK_GROUP_ITEM=192
 	Used          int64           `bin:"off=0, siz=8"`
 	ChunkObjectID internal.ObjID  `bin:"off=8, siz=8"`
diff --git a/pkg/btrfs/btrfsitem/item_chunk.go b/pkg/btrfs/btrfsitem/item_chunk.go
index eb3a16d..8f98ba1 100644
--- a/pkg/btrfs/btrfsitem/item_chunk.go
+++ b/pkg/btrfs/btrfsitem/item_chunk.go
@@ -8,21 +8,24 @@ import (
 )
 
 // Maps logical address to physical.
+//
+// key.objectid = BTRFS_FIRST_CHUNK_TREE_OBJECTID
+// key.offset = logical_addr
 type Chunk struct { // CHUNK_ITEM=228
 	Head    ChunkHeader
 	Stripes []ChunkStripe
 }
 
 type ChunkHeader struct {
-	Size           uint64          `bin:"off=0x0,  siz=0x8"`
-	Owner          internal.ObjID  `bin:"off=0x8,  siz=0x8"` // root referencing this chunk (always EXTENT_TREE_OBJECTID=2)
-	StripeLen      uint64          `bin:"off=0x10, siz=0x8"` // ???
-	Type           BlockGroupFlags `bin:"off=0x18, siz=0x8"`
-	IOOptimalAlign uint32          `bin:"off=0x20, siz=0x4"`
-	IOOptimalWidth uint32          `bin:"off=0x24, siz=0x4"`
-	IOMinSize      uint32          `bin:"off=0x28, siz=0x4"` // sector size
-	NumStripes     uint16          `bin:"off=0x2c, siz=0x2"` // [ignored-when-writing]
-	SubStripes     uint16          `bin:"off=0x2e, siz=0x2"` // ???
+	Size           internal.AddrDelta `bin:"off=0x0,  siz=0x8"`
+	Owner          internal.ObjID     `bin:"off=0x8,  siz=0x8"` // root referencing this chunk (always EXTENT_TREE_OBJECTID=2)
+	StripeLen      uint64             `bin:"off=0x10, siz=0x8"` // ???
+	Type           BlockGroupFlags    `bin:"off=0x18, siz=0x8"`
+	IOOptimalAlign uint32             `bin:"off=0x20, siz=0x4"`
+	IOOptimalWidth uint32             `bin:"off=0x24, siz=0x4"`
+	IOMinSize      uint32             `bin:"off=0x28, siz=0x4"` // sector size
+	NumStripes     uint16             `bin:"off=0x2c, siz=0x2"` // [ignored-when-writing]
+	SubStripes     uint16             `bin:"off=0x2e, siz=0x2"` // ???
 	binstruct.End  `bin:"off=0x30"`
 }
 
diff --git a/pkg/btrfs/btrfsitem/item_devextent.go b/pkg/btrfs/btrfsitem/item_devextent.go
index dfa6d03..7f08b5e 100644
--- a/pkg/btrfs/btrfsitem/item_devextent.go
+++ b/pkg/btrfs/btrfsitem/item_devextent.go
@@ -5,11 +5,13 @@ import (
 	"lukeshu.com/btrfs-tools/pkg/btrfs/internal"
 )
 
+// key.objectid = device_id
+// key.offset = physical_addr
 type DevExtent struct { // DEV_EXTENT=204
 	ChunkTree     int64                `bin:"off=0, siz=8"`
 	ChunkObjectID internal.ObjID       `bin:"off=8, siz=8"`
 	ChunkOffset   internal.LogicalAddr `bin:"off=16, siz=8"`
-	Length        uint64               `bin:"off=24, siz=8"`
+	Length        internal.AddrDelta   `bin:"off=24, siz=8"`
 	ChunkTreeUUID internal.UUID        `bin:"off=32, siz=16"`
 	binstruct.End `bin:"off=48"`
 }
diff --git a/pkg/btrfs/internal.go b/pkg/btrfs/internal.go
index 5c0fc8b..e00e946 100644
--- a/pkg/btrfs/internal.go
+++ b/pkg/btrfs/internal.go
@@ -11,6 +11,7 @@ type (
 	ObjID        = internal.ObjID
 	LogicalAddr  = internal.LogicalAddr
 	PhysicalAddr = internal.PhysicalAddr
+	AddrDelta    = internal.AddrDelta
 
 	// complex types
 
diff --git a/pkg/btrfs/internal/addr.go b/pkg/btrfs/internal/addr.go
index 7067982..a3f41ed 100644
--- a/pkg/btrfs/internal/addr.go
+++ b/pkg/btrfs/internal/addr.go
@@ -9,6 +9,7 @@ import (
 type (
 	PhysicalAddr int64
 	LogicalAddr  int64
+	AddrDelta    int64
 )
 
 func formatAddr(addr int64, f fmt.State, verb rune) {
@@ -23,3 +24,10 @@ func formatAddr(addr int64, f fmt.State, verb rune) {
 
 func (a PhysicalAddr) Format(f fmt.State, verb rune) { formatAddr(int64(a), f, verb) }
 func (a LogicalAddr) Format(f fmt.State, verb rune)  { formatAddr(int64(a), f, verb) }
+func (d AddrDelta) Format(f fmt.State, verb rune)    { formatAddr(int64(d), f, verb) }
+
+func (a PhysicalAddr) Sub(b PhysicalAddr) AddrDelta { return AddrDelta(a - b) }
+func (a LogicalAddr) Sub(b LogicalAddr) AddrDelta   { return AddrDelta(a - b) }
+
+func (a PhysicalAddr) Add(b AddrDelta) PhysicalAddr { return a + PhysicalAddr(b) }
+func (a LogicalAddr) Add(b AddrDelta) LogicalAddr   { return a + LogicalAddr(b) }
diff --git a/pkg/btrfs/io2_fs.go b/pkg/btrfs/io2_fs.go
index 588a60a..8b76e2d 100644
--- a/pkg/btrfs/io2_fs.go
+++ b/pkg/btrfs/io2_fs.go
@@ -150,20 +150,20 @@ type QualifiedPhysicalAddr struct {
 	Addr PhysicalAddr
 }
 
-func (fs *FS) Resolve(laddr LogicalAddr) (paddrs map[QualifiedPhysicalAddr]struct{}, maxlen uint64) {
+func (fs *FS) Resolve(laddr LogicalAddr) (paddrs map[QualifiedPhysicalAddr]struct{}, maxlen AddrDelta) {
 	paddrs = make(map[QualifiedPhysicalAddr]struct{})
-	maxlen = math.MaxUint64
+	maxlen = math.MaxInt64
 
 	for _, chunk := range fs.chunks {
 		low := LogicalAddr(chunk.Key.Offset)
-		high := low + LogicalAddr(chunk.Chunk.Head.Size)
+		high := low.Add(chunk.Chunk.Head.Size)
 		if low <= laddr && laddr < high {
-			offsetWithinChunk := uint64(laddr) - chunk.Key.Offset
+			offsetWithinChunk := laddr.Sub(low)
 			maxlen = util.Min(maxlen, chunk.Chunk.Head.Size-offsetWithinChunk)
 			for _, stripe := range chunk.Chunk.Stripes {
 				paddrs[QualifiedPhysicalAddr{
 					Dev:  stripe.DeviceUUID,
-					Addr: stripe.Offset + PhysicalAddr(offsetWithinChunk),
+					Addr: stripe.Offset.Add(offsetWithinChunk),
 				}] = struct{}{}
 			}
 		}
@@ -189,7 +189,7 @@ func (fs *FS) maybeShortReadAt(dat []byte, laddr LogicalAddr) (int, error) {
 	if len(paddrs) == 0 {
 		return 0, fmt.Errorf("read: could not map logical address %v", laddr)
 	}
-	if uint64(len(dat)) > maxlen {
+	if AddrDelta(len(dat)) > maxlen {
 		dat = dat[:maxlen]
 	}
 
@@ -231,7 +231,7 @@ func (fs *FS) maybeShortWriteAt(dat []byte, laddr LogicalAddr) (int, error) {
 	if len(paddrs) == 0 {
 		return 0, fmt.Errorf("write: could not map logical address %v", laddr)
 	}
-	if uint64(len(dat)) > maxlen {
+	if AddrDelta(len(dat)) > maxlen {
 		dat = dat[:maxlen]
 	}