From 332aee670ca7f76f6f0fc07eb1bdb781cc7a8ffe Mon Sep 17 00:00:00 2001
From: Scott Moser <scmoser@cisco.com>
Date: Wed, 8 Apr 2020 16:04:11 -0400
Subject: [PATCH] Add range checking in addPartitionSet.

This checks that partitions provided to addPartitionSet are inside
the ranges allowed.

Note that it does still allow for partitions to overlap themselves
or to overlap existing partitions, but will give a reasonable error
if you try to give a partition.Last or partition.New that are not sane.
---
 linux/disk.go      | 15 ++++++++++++-
 linux/disk_test.go | 52 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/linux/disk.go b/linux/disk.go
index 8e973b0..bea1436 100644
--- a/linux/disk.go
+++ b/linux/disk.go
@@ -351,11 +351,24 @@ func addPartitionSet(d disko.Disk, pSet disko.PartitionSet) error {
 		return err
 	}
 
+	maxEnd := ((d.Size - uint64(d.SectorSize)*33) / disko.Mebibyte) * disko.Mebibyte
+	minStart := disko.Mebibyte
+
 	for _, p := range pSet {
 		gptTable.Partitions[p.Number-1] = toGPTPartition(p, d.SectorSize)
 
+		if p.Start < minStart {
+			return fmt.Errorf("partition %d start (%d) is too low. Must be >= %d",
+				p.Number, p.Start, minStart)
+		}
+
+		if p.Last >= maxEnd {
+			return fmt.Errorf("partition %d Last (%d) is too high. Must be < %d",
+				p.Number, p.Last, maxEnd)
+		}
+
 		if err := zeroStartEnd(fp, int64(p.Start), int64(p.Last)); err != nil {
-			return fmt.Errorf("failed to zero partition %d: %s", p.ID, err)
+			return fmt.Errorf("failed to zero partition %d: %s", p.Number, err)
 		}
 	}
 
diff --git a/linux/disk_test.go b/linux/disk_test.go
index 7d02101..2163485 100644
--- a/linux/disk_test.go
+++ b/linux/disk_test.go
@@ -216,3 +216,55 @@ func TestDeletePartition(t *testing.T) {
 		t.Fatalf("There were %d partitions after delete, expected 0", len(pSet))
 	}
 }
+
+func TestBadPartition(t *testing.T) {
+	tmpd, err := ioutil.TempDir("", "disko_test")
+	if err != nil {
+		t.Fatalf("Failed to create tempdir: %s", err)
+	}
+
+	defer os.RemoveAll(tmpd)
+
+	fpath := path.Join(tmpd, "mydisk")
+	fsize := uint64(200 * 1024 * 1024) // nolint:gomnd
+
+	if err := ioutil.WriteFile(fpath, []byte{}, 0644); err != nil {
+		t.Fatalf("Failed to write to a temp file: %s", err)
+	}
+
+	if err := os.Truncate(fpath, int64(fsize)); err != nil {
+		t.Fatalf("Failed create empty file: %s", err)
+	}
+
+	disk := disko.Disk{
+		Name:       "mydisk",
+		Path:       fpath,
+		Size:       fsize,
+		SectorSize: sectorSize512,
+	}
+
+	fs := disk.FreeSpaces()
+	myGUID := disko.GenGUID()
+
+	part := disko.Partition{
+		Start:  1024, // nolint:gomnd
+		Last:   fs[0].Last,
+		Type:   partid.LinuxLVM,
+		Name:   "mytest partition",
+		ID:     myGUID,
+		Number: uint(1),
+	}
+
+	err = addPartitionSet(disk, disko.PartitionSet{part.Number: part})
+	if err == nil {
+		t.Errorf("Created partition with OOB start (%d). should have failed", part.Start)
+	}
+
+	part.Start = fs[0].Start
+	part.Last = disk.Size - 1
+
+	err = addPartitionSet(disk, disko.PartitionSet{part.Number: part})
+	if err == nil {
+		t.Errorf("Created partition with OOB end (%d). should have failed", part.Last)
+	}
+}