seal_oem_partition: make cmdline updates compatible with pickargs

Update behavior to not assume device mapper table is at the end of the
command line.

BUG=b/517096251

Change-Id: I9842c48689127dac7f8aad2a87a30354fd672250
Reviewed-on: https://cos-review.googlesource.com/c/cos/tools/+/156144
Tested-by: Robert Kolchmeyer <rkolchmeyer@google.com>
Reviewed-by: He Gao <hegao@google.com>
Cloud-Build: 228075978874@cloudbuild.gserviceaccount.com <228075978874@cloudbuild.gserviceaccount.com>
Reviewed-by: Miri Amarilio <mirilio@google.com>
diff --git a/src/pkg/tools/BUILD.bazel b/src/pkg/tools/BUILD.bazel
index 6af499d..a25ade3 100644
--- a/src/pkg/tools/BUILD.bazel
+++ b/src/pkg/tools/BUILD.bazel
@@ -29,7 +29,10 @@
 
 go_test(
     name = "tools_test",
-    srcs = ["handle_disk_layout_test.go"],
+    srcs = [
+        "handle_disk_layout_test.go",
+        "seal_oem_partition_test.go",
+    ],
     embed = [":tools"],
     deps = ["//src/pkg/tools/partutil/partutiltest"],
 )
diff --git a/src/pkg/tools/seal_oem_partition.go b/src/pkg/tools/seal_oem_partition.go
index 5a36492..b0fee2c 100644
--- a/src/pkg/tools/seal_oem_partition.go
+++ b/src/pkg/tools/seal_oem_partition.go
@@ -22,6 +22,7 @@
 	"os"
 	"os/exec"
 	"path/filepath"
+	"regexp"
 	"runtime"
 	"strconv"
 	"strings"
@@ -29,6 +30,11 @@
 	"cos.googlesource.com/cos/tools.git/src/pkg/tools/partutil"
 )
 
+var (
+	dmV0Re = regexp.MustCompile(`\bdm="([0-9])([^"]*)"`)
+	dmV1Re = regexp.MustCompile(`\bdm-mod\.create="([^"]*)"`)
+)
+
 // SealOEMPartition sets the hashtree of the OEM partition
 // with "veritysetup" and modifies the kernel command line to
 // verify the OEM partition at boot time.
@@ -201,41 +207,32 @@
 	// from 4K blocks to 512B sectors
 	oemFSSizeSector := oemFSSize4K << 3
 	entryStringV0 := fmt.Sprintf("%s none ro 1, 0 %d verity payload=PARTUUID=%s hashtree=PARTUUID=%s "+
-		"hashstart=%d alg=sha256 root_hexdigest=%s salt=%s\"", name, oemFSSizeSector,
+		"hashstart=%d alg=sha256 root_hexdigest=%s salt=%s", name, oemFSSizeSector,
 		partUUID, partUUID, oemFSSizeSector, hash, salt)
 	entryStringV1 := fmt.Sprintf("%s,,,ro,0 %d verity 0 PARTUUID=%s PARTUUID=%s "+
-		"4096 4096 %d %d sha256 %s %s\"", name, oemFSSizeSector,
+		"4096 4096 %d %d sha256 %s %s", name, oemFSSizeSector,
 		partUUID, partUUID, oemFSSize4K, oemFSSize4K, hash, salt)
+
+	// Escape '$' in entry strings for safe use in regex replacement templates.
+	entryStringV0Escaped := strings.ReplaceAll(entryStringV0, "$", "$$")
+	entryStringV1Escaped := strings.ReplaceAll(entryStringV1, "$", "$$")
+
 	grubContent, err := ioutil.ReadFile(grubPath)
 	if err != nil {
 		return fmt.Errorf("cannot read grub.cfg at %q, "+
 			"input: grubPath=%q, name=%q, partUUID=%q, oemFSSize4K=%d, hash=%q, salt=%q, "+
 			"error msg:(%v)", grubPath, grubPath, name, partUUID, oemFSSize4K, hash, salt, err)
 	}
+
 	lines := strings.Split(string(grubContent), "\n")
-	// add the entry to all kernel command lines containing "dm="
-	dmVersion := 0
 	for idx, line := range lines {
-		if !strings.Contains(line, "dm=") &&
-			!strings.Contains(line, "dm-mod.create=") {
-			continue
-		}
-		var startPos = strings.Index(line, "dm=")
-		if startPos == -1 {
-			startPos = strings.Index(line, "dm-mod.create=")
-			dmVersion = 1
-		}
-		// remove the end quote.
-		lineBuf := []rune(line[:len(line)-1])
-		if dmVersion == 0 {
-			// add number of entries.
-			lineBuf[startPos+4] = '2'
-			lines[idx] = strings.Join(append(strings.Split(string(lineBuf), ","), entryStringV0), ",")
-		} else {
-			configs := []string{string(lineBuf), entryStringV1}
-			lines[idx] = strings.Join(configs, ";")
+		if strings.Contains(line, "dm=") {
+			lines[idx] = dmV0Re.ReplaceAllString(line, fmt.Sprintf("dm=\"2${2},%s\"", entryStringV0Escaped))
+		} else if strings.Contains(line, "dm-mod.create=") {
+			lines[idx] = dmV1Re.ReplaceAllString(line, fmt.Sprintf("dm-mod.create=\"${1};%s\"", entryStringV1Escaped))
 		}
 	}
+
 	// new content of grub.cfg
 	grubContent = []byte(strings.Join(lines, "\n"))
 	err = ioutil.WriteFile(grubPath, grubContent, 0755)
diff --git a/src/pkg/tools/seal_oem_partition_test.go b/src/pkg/tools/seal_oem_partition_test.go
new file mode 100644
index 0000000..c489cf3
--- /dev/null
+++ b/src/pkg/tools/seal_oem_partition_test.go
@@ -0,0 +1,149 @@
+// Copyright 2026 Google LLC
+//
+// 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 tools
+
+import (
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+func TestAppendDMEntryToGRUB(t *testing.T) {
+	tests := []struct {
+		name          string
+		initialConfig string
+		grubName      string
+		partUUID      string
+		hash          string
+		salt          string
+		oemFSSize4K   uint64
+		wantConfig    string
+		wantErr       bool
+	}{
+		{
+			name: "V0 (dm=)",
+			initialConfig: `menuentry "COS" {
+	linux /boot/vmlinuz root=/dev/dm-0 dm="1 vroot none ro 1,0 4077568 verity payload=PARTUUID=1234 hashtree=PARTUUID=1234 hashstart=4077568 alg=sha256 root_hexdigest=hash0 salt=salt0"
+}`,
+			grubName:    "oemroot",
+			partUUID:    "oem-uuid",
+			hash:        "oem-hash",
+			salt:        "oem-salt",
+			oemFSSize4K: 1000,
+			wantConfig: `menuentry "COS" {
+	linux /boot/vmlinuz root=/dev/dm-0 dm="2 vroot none ro 1,0 4077568 verity payload=PARTUUID=1234 hashtree=PARTUUID=1234 hashstart=4077568 alg=sha256 root_hexdigest=hash0 salt=salt0,oemroot none ro 1, 0 8000 verity payload=PARTUUID=oem-uuid hashtree=PARTUUID=oem-uuid hashstart=8000 alg=sha256 root_hexdigest=oem-hash salt=oem-salt"
+}`,
+			wantErr: false,
+		},
+		{
+			name: "V1 (dm-mod.create=)",
+			initialConfig: `menuentry "COS" {
+	linux /boot/vmlinuz root=/dev/dm-0 dm-mod.create="vroot,,,ro,0 4077568 verity 0 PARTUUID=1234 PARTUUID=1234 4096 4096 509696 509696 sha256 hash0 salt0"
+}`,
+			grubName:    "oemroot",
+			partUUID:    "oem-uuid",
+			hash:        "oem-hash",
+			salt:        "oem-salt",
+			oemFSSize4K: 1000,
+			wantConfig: `menuentry "COS" {
+	linux /boot/vmlinuz root=/dev/dm-0 dm-mod.create="vroot,,,ro,0 4077568 verity 0 PARTUUID=1234 PARTUUID=1234 4096 4096 509696 509696 sha256 hash0 salt0;oemroot,,,ro,0 8000 verity 0 PARTUUID=oem-uuid PARTUUID=oem-uuid 4096 4096 1000 1000 sha256 oem-hash oem-salt"
+}`,
+			wantErr: false,
+		},
+		{
+			name: "No match",
+			initialConfig: `menuentry "COS" {
+	linux /boot/vmlinuz root=/dev/sda3
+}`,
+			grubName:    "oemroot",
+			partUUID:    "oem-uuid",
+			hash:        "oem-hash",
+			salt:        "oem-salt",
+			oemFSSize4K: 1000,
+			wantConfig: `menuentry "COS" {
+	linux /boot/vmlinuz root=/dev/sda3
+}`,
+			wantErr: false,
+		},
+		{
+			name: "Multiple matches",
+			initialConfig: `menuentry "COS A" {
+	linux /boot/vmlinuz root=/dev/dm-0 dm="1 vroot none ro 1,0 4077568 verity payload=PARTUUID=1234 hashtree=PARTUUID=1234 hashstart=4077568 alg=sha256 root_hexdigest=hash0 salt=salt0"
+}
+menuentry "COS B" {
+	linux /boot/vmlinuz root=/dev/dm-0 dm="1 vroot none ro 1,0 4077568 verity payload=PARTUUID=1234 hashtree=PARTUUID=1234 hashstart=4077568 alg=sha256 root_hexdigest=hash0 salt=salt0"
+}`,
+			grubName:    "oemroot",
+			partUUID:    "oem-uuid",
+			hash:        "oem-hash",
+			salt:        "oem-salt",
+			oemFSSize4K: 1000,
+			wantConfig: `menuentry "COS A" {
+	linux /boot/vmlinuz root=/dev/dm-0 dm="2 vroot none ro 1,0 4077568 verity payload=PARTUUID=1234 hashtree=PARTUUID=1234 hashstart=4077568 alg=sha256 root_hexdigest=hash0 salt=salt0,oemroot none ro 1, 0 8000 verity payload=PARTUUID=oem-uuid hashtree=PARTUUID=oem-uuid hashstart=8000 alg=sha256 root_hexdigest=oem-hash salt=oem-salt"
+}
+menuentry "COS B" {
+	linux /boot/vmlinuz root=/dev/dm-0 dm="2 vroot none ro 1,0 4077568 verity payload=PARTUUID=1234 hashtree=PARTUUID=1234 hashstart=4077568 alg=sha256 root_hexdigest=hash0 salt=salt0,oemroot none ro 1, 0 8000 verity payload=PARTUUID=oem-uuid hashtree=PARTUUID=oem-uuid hashstart=8000 alg=sha256 root_hexdigest=oem-hash salt=oem-salt"
+}`,
+			wantErr: false,
+		},
+		{
+			name: "Pickargs cmdline",
+			initialConfig: `menuentry "COS" {
+  linux /boot/vmlinuz root=/dev/dm-0 dm-mod.create="vroot,,,ro,0 4077568 verity 0 PARTUUID=1234 PARTUUID=1234 4096 4096 509696 509696 sha256 hash0 salt0" $cmdline_extra
+}`,
+			grubName:    "oemroot",
+			partUUID:    "oem-uuid",
+			hash:        "oem-hash",
+			salt:        "oem-salt",
+			oemFSSize4K: 1000,
+			wantConfig: `menuentry "COS" {
+  linux /boot/vmlinuz root=/dev/dm-0 dm-mod.create="vroot,,,ro,0 4077568 verity 0 PARTUUID=1234 PARTUUID=1234 4096 4096 509696 509696 sha256 hash0 salt0;oemroot,,,ro,0 8000 verity 0 PARTUUID=oem-uuid PARTUUID=oem-uuid 4096 4096 1000 1000 sha256 oem-hash oem-salt" $cmdline_extra
+}`,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tmpDir, err := ioutil.TempDir("", "grub-test")
+			if err != nil {
+				t.Fatalf("failed to create temp dir: %v", err)
+			}
+			defer os.RemoveAll(tmpDir)
+
+			grubPath := filepath.Join(tmpDir, "grub.cfg")
+			err = ioutil.WriteFile(grubPath, []byte(tt.initialConfig), 0644)
+			if err != nil {
+				t.Fatalf("failed to write initial grub config: %v", err)
+			}
+
+			err = appendDMEntryToGRUB(grubPath, tt.grubName, tt.partUUID, tt.hash, tt.salt, tt.oemFSSize4K)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("appendDMEntryToGRUB() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+
+			if !tt.wantErr {
+				gotConfig, err := ioutil.ReadFile(grubPath)
+				if err != nil {
+					t.Fatalf("failed to read modified grub config: %v", err)
+				}
+				if string(gotConfig) != tt.wantConfig {
+					t.Errorf("appendDMEntryToGRUB() got =\n%s\nwant =\n%s", string(gotConfig), tt.wantConfig)
+				}
+			}
+		})
+	}
+}