OSModifier: allow two linux cmdline in grub.cfg (#10472)

Author: elainezhao1

toolkit/tools/pkg/imagecustomizerlib/grubcfgutils.go

@@ -128,6 +128,50 @@ func FindLinuxLine(inputGrubCfgContent string) (grub.Line, error) {
 	return lines[0], nil
 }
 
+// Find the linux command within non-recovery mode menuentry block in the grub config file.
+func FindNonRecoveryLinuxLine(inputGrubCfgContent string) ([]grub.Line, error) {
+	grubTokens, err := grub.TokenizeConfig(inputGrubCfgContent)
+	if err != nil {
+		return nil, err
+	}
+
+	grubLines := grub.SplitTokensIntoLines(grubTokens)
+	var linuxLines []grub.Line
+	inMenuEntry := false
+	isRecoveryMenu := false
+
+	// Iterate over all lines to find non-recovery mode menuentry and its linux line
+	for _, line := range grubLines {
+		if len(line.Tokens) > 1 && grub.IsTokenKeyword(line.Tokens[0], "menuentry") {
+			// Found a new 'menuentry', reset flags
+			inMenuEntry = true
+			isRecoveryMenu = false
+
+			// Check if the title (second token) contains the word 'recovery'
+			if strings.Contains(line.Tokens[1].RawContent, "recovery") {
+				isRecoveryMenu = true
+			}
+
+			// If it's a recovery menuentry, ignore this block
+			if isRecoveryMenu {
+				inMenuEntry = false
+			}
+		} else if inMenuEntry {
+			// We are inside a non-recovery menuentry block
+			if len(line.Tokens) > 0 && grub.IsTokenKeyword(line.Tokens[0], "linux") {
+				// Append only lines that contain the 'linux' command
+				linuxLines = append(linuxLines, line)
+			}
+		}
+	}
+
+	if len(linuxLines) == 0 {
+		return nil, fmt.Errorf("no linux line found in non-recovery menuentry")
+	}
+
+	return linuxLines, nil
+}
+
 // Overrides the path of the kernel binary of all the linux commands within a grub config file.
 func setLinuxOrInitrdPathAll(inputGrubCfgContent string, commandName string, filePath string, allowMultiple bool) (outputGrubCfgContent string, oldFilePaths []string, err error) {
 	quotedFilePath := grub.QuoteString(filePath)

toolkit/tools/pkg/osmodifierlib/modifydefaultgrub.go

@@ -64,12 +64,16 @@ func extractValuesFromGrubConfig(imageChroot safechroot.ChrootInterface) ([]stri
 		return nil, err
 	}
 
-	line, err := imagecustomizerlib.FindLinuxLine(grubCfgContent)
+	lines, err := imagecustomizerlib.FindNonRecoveryLinuxLine(grubCfgContent)
 	if err != nil {
 		return nil, err
 	}
 
-	argTokens, err := imagecustomizerlib.ParseCommandLineArgs(line.Tokens)
+	if len(lines) != 1 {
+		return nil, fmt.Errorf("expected 1 non-recovery linux line, found %d", len(lines))
+	}
+
+	argTokens, err := imagecustomizerlib.ParseCommandLineArgs(lines[0].Tokens)
 	if err != nil {
 		return nil, err
 	}