From 1434a1ab1ba7c5ea0ae17530183cf44687cab324 Mon Sep 17 00:00:00 2001
From: Jonathan Lebon <jonathan@jlebon.com>
Date: Wed, 25 Oct 2023 11:57:01 -0400
Subject: [PATCH] stages/disks: retry `sgdisk --zap-all` invocation

When `wipeTable` is enabled, we run `sgdisk --zap-all`. But if the table
was corrupted in the first place, `sgdisk` will exit with code 2 which
then breaks boot.

As a workaround, Ignition used to retry the invocation but the context
around it was lost in #544 and #1149 and the retry was removed and
the error-checking was added.

So this patch effectively re-applies 94c98bcb ("sgdisk: retry zap-all
operation on failure"), but now with a comment and a test to make sure
we don't regress again.

Closes: https://github.com/coreos/fedora-coreos-tracker/issues/1596
---
 docs/release-notes.md                    |  1 +
 internal/exec/stages/disks/partitions.go |  7 ++-
 tests/filesystem.go                      | 15 ++++++
 tests/positive/partitions/wipe.go        | 63 ++++++++++++++++++++++++
 tests/types/types.go                     |  9 ++--
 5 files changed, 90 insertions(+), 5 deletions(-)
 create mode 100644 tests/positive/partitions/wipe.go

diff --git a/docs/release-notes.md b/docs/release-notes.md
index 3411ac0624..165d5bfc90 100644
--- a/docs/release-notes.md
+++ b/docs/release-notes.md
@@ -21,6 +21,7 @@ nav_order: 9
 ### Bug fixes
 
 - Prevent races with udev after disk editing
+- Don't fail to wipe partition table if it's corrupted
 
 
 ## Ignition 2.16.2 (2023-07-12)
diff --git a/internal/exec/stages/disks/partitions.go b/internal/exec/stages/disks/partitions.go
index a21fc0477c..cb55c37655 100644
--- a/internal/exec/stages/disks/partitions.go
+++ b/internal/exec/stages/disks/partitions.go
@@ -324,7 +324,12 @@ func (s stage) partitionDisk(dev types.Disk, devAlias string) error {
 		s.Logger.Info("wiping partition table requested on %q", devAlias)
 		op.WipeTable(true)
 		if err := op.Commit(); err != nil {
-			return err
+			// `sgdisk --zap-all` will exit code 2 if the table was corrupted; retry it
+			// https://github.com/coreos/fedora-coreos-tracker/issues/1596
+			s.Logger.Info("potential error encountered while wiping table... retrying")
+			if err := op.Commit(); err != nil {
+				return err
+			}
 		}
 	}
 
diff --git a/tests/filesystem.go b/tests/filesystem.go
index fa16c44b97..f2ae938cfc 100644
--- a/tests/filesystem.go
+++ b/tests/filesystem.go
@@ -19,6 +19,7 @@ import (
 	"bytes"
 	"compress/bzip2"
 	"context"
+	"crypto/rand"
 	"encoding/base64"
 	"fmt"
 	"io"
@@ -226,6 +227,20 @@ func setupDisk(ctx context.Context, disk *types.Disk, diskIndex int, imageSize i
 			return err
 		}
 	}
+
+	if disk.CorruptTable {
+		bytes := make([]byte, 1536)
+		if _, err := rand.Read(bytes); err != nil {
+			return err
+		}
+		f, err := os.OpenFile(disk.ImageFile, os.O_WRONLY, 0666)
+		if err != nil {
+			return err
+		}
+		f.Write(bytes)
+		defer f.Close()
+	}
+
 	return nil
 }
 
diff --git a/tests/positive/partitions/wipe.go b/tests/positive/partitions/wipe.go
new file mode 100644
index 0000000000..d51b5d59fb
--- /dev/null
+++ b/tests/positive/partitions/wipe.go
@@ -0,0 +1,63 @@
+// Copyright 2023 CoreOS, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package partitions
+
+import (
+	"github.com/coreos/ignition/v2/tests/register"
+	"github.com/coreos/ignition/v2/tests/types"
+)
+
+func init() {
+	// Tests that wipe tables
+	register.Register(register.PositiveTest, WipeBadTable())
+}
+
+func WipeBadTable() types.Test {
+	name := "partition.wipebadtable"
+	in := append(types.GetBaseDisk(), types.Disk{
+		Alignment: types.IgnitionAlignment,
+		Partitions: types.Partitions{
+			{
+				Label:  "deleteme",
+				Number: 1,
+				Length: 65536,
+			},
+		},
+		CorruptTable: true,
+	})
+	out := append(types.GetBaseDisk(), types.Disk{Alignment: types.IgnitionAlignment})
+	config := `{
+		"ignition": {
+			"version": "$version"
+		},
+		"storage": {
+			"disks": [
+			{
+				"device": "$disk1",
+				"wipeTable": true
+			}
+			]
+		}
+	}`
+	configMinVersion := "3.0.0"
+
+	return types.Test{
+		Name:             name,
+		In:               in,
+		Out:              out,
+		Config:           config,
+		ConfigMinVersion: configMinVersion,
+	}
+}
diff --git a/tests/types/types.go b/tests/types/types.go
index d5bd3fe78a..4405cb6f08 100644
--- a/tests/types/types.go
+++ b/tests/types/types.go
@@ -57,10 +57,11 @@ type Node struct {
 }
 
 type Disk struct {
-	ImageFile  string
-	Device     string
-	Alignment  int
-	Partitions Partitions
+	ImageFile    string
+	Device       string
+	Alignment    int
+	Partitions   Partitions
+	CorruptTable bool // set to true to corrupt the partition table
 }
 
 type Partitions []*Partition