Fix args preparation for re-execution in boot and gofer process.

In runsc, the boot (runsc/cmd/boot.go) process and gofer (runsc/cmd/gofer.go)
process re-execute themselves (for various reasons). For re-execution, they
modify some flags in os.Args. There were certain issues with the process of
generating the args for re-execution that this patch aims to fix:
- In the gofer, we were not overwriting the flag if there was a conflict. This
  could lead to duplicate flags in the args with conflicting values.
- The arg manipulation (see old Boot.prepareArgs()) was kinda hacky in figuring
  out where to place new arguments.
- Arg manipulation was using `strings.Contains()` to figure out which flag to
  replace. This can be problematic when manipulating flags that are substrings
  of other flags (like --profile and --profile-mutex). This can also hurt when
  random flag values contain the flag name string.
- There was no common general utility to do this work.

This patch adds a generic `prepareArgs()` function to the cmd package which
can be used by boot and gofer and it fixes all the above-mentioned issues.

PiperOrigin-RevId: 552627805

runsc/cmd/boot.go

@@ -224,6 +224,7 @@ func (b *Boot) Execute(_ context.Context, f *flag.FlagSet, args ...any) subcomma
 	// Initialize ring0 library.
 	ring0.InitDefault()
 
+	argOverride := make(map[string]string)
 	if len(b.productName) == 0 {
 		// Do this before chroot takes effect, otherwise we can't read /sys.
 		if product, err := ioutil.ReadFile("/sys/devices/virtual/dmi/id/product_name"); err != nil {
@@ -231,6 +232,7 @@ func (b *Boot) Execute(_ context.Context, f *flag.FlagSet, args ...any) subcomma
 		} else {
 			b.productName = strings.TrimSpace(string(product))
 			log.Infof("Setting product_name: %q", b.productName)
+			argOverride["product-name"] = b.productName
 		}
 	}
 
@@ -243,7 +245,10 @@ func (b *Boot) Execute(_ context.Context, f *flag.FlagSet, args ...any) subcomma
 		}
 	}
 
-	syncUsernsForRootless(b.syncUsernsFD)
+	if b.syncUsernsFD >= 0 {
+		syncUsernsForRootless(b.syncUsernsFD)
+		argOverride["sync-userns-fd"] = "-1"
+	}
 
 	// Get the spec from the specFD. We *must* keep this os.File alive past
 	// the call setCapsAndCallSelf, otherwise the FD will be closed and the
@@ -258,6 +263,7 @@ func (b *Boot) Execute(_ context.Context, f *flag.FlagSet, args ...any) subcomma
 		if err := setUpChroot(b.pidns, spec, conf); err != nil {
 			util.Fatalf("error setting up chroot: %v", err)
 		}
+		argOverride["setup-root"] = "false"
 
 		if !conf.Rootless {
 			// /proc is umounted from a forked process, because the
@@ -270,6 +276,7 @@ func (b *Boot) Execute(_ context.Context, f *flag.FlagSet, args ...any) subcomma
 				panic("procMountSyncFD is set")
 			}
 			b.procMountSyncFD = int(w.Fd())
+			argOverride["proc-mount-sync-fd"] = strconv.Itoa(b.procMountSyncFD)
 
 			// Clear FD_CLOEXEC. Regardless of b.applyCaps, this process will be
 			// re-executed. procMountSyncFD should remain open.
@@ -279,7 +286,7 @@ func (b *Boot) Execute(_ context.Context, f *flag.FlagSet, args ...any) subcomma
 
 			if !b.applyCaps {
 				// Remove the args that have already been done before calling self.
-				args := b.prepareArgs("setup-root", "sync-userns-fd")
+				args := prepareArgs(b.Name(), f, argOverride)
 
 				// Note that we've already read the spec from the spec FD, and
 				// we will read it again after the exec call. This works
@@ -314,9 +321,10 @@ func (b *Boot) Execute(_ context.Context, f *flag.FlagSet, args ...any) subcomma
 		if conf.DirectFS {
 			caps = specutils.MergeCapabilities(caps, directfsSandboxLinuxCaps)
 		}
+		argOverride["apply-caps"] = "false"
 
 		// Remove the args that have already been done before calling self.
-		args := b.prepareArgs("setup-root", "sync-userns-fd", "apply-caps")
+		args := prepareArgs(b.Name(), f, argOverride)
 
 		// Note that we've already read the spec from the spec FD, and
 		// we will read it again after the exec call. This works
@@ -466,28 +474,38 @@ func (b *Boot) Execute(_ context.Context, f *flag.FlagSet, args ...any) subcomma
 	return subcommands.ExitSuccess
 }
 
-func (b *Boot) prepareArgs(exclude ...string) []string {
+// prepareArgs returns the args that can be used to re-execute the current
+// program. It manipulates the flags of the subcommands.Command identified by
+// subCmdName and fSet is the flag.FlagSet of this subcommand. It applies the
+// flags specified by override map. In case of conflict, flag is overriden.
+//
+// Postcondition: prepareArgs() takes ownership of override map.
+func prepareArgs(subCmdName string, fSet *flag.FlagSet, override map[string]string) []string {
 	var args []string
+	// Add all args up until (and including) the sub command.
 	for _, arg := range os.Args {
-		for _, excl := range exclude {
-			if strings.Contains(arg, excl) {
-				goto skip
-			}
-		}
 		args = append(args, arg)
-		// Some parameters are not already part of os.Args because they are
-		// solely configured by Boot.Execute(). Strategically add these parameters
-		// after the command and before the container ID at the end.
-		if arg == "boot" {
-			if b.procMountSyncFD != -1 {
-				args = append(args, fmt.Sprintf("--proc-mount-sync-fd=%d", b.procMountSyncFD))
-			}
-			if len(b.productName) > 0 {
-				args = append(args, "--product-name", b.productName)
-			}
+		if arg == subCmdName {
+			break
 		}
-	skip:
 	}
+	// Set sub command flags. Iterate through all the explicitly set flags.
+	fSet.Visit(func(gf *flag.Flag) {
+		// If a conflict is found with override, then prefer override flag.
+		if ov, ok := override[gf.Name]; ok {
+			args = append(args, fmt.Sprintf("--%s=%s", gf.Name, ov))
+			delete(override, gf.Name)
+			return
+		}
+		// Otherwise pass through the original flag.
+		args = append(args, fmt.Sprintf("--%s=%s", gf.Name, gf.Value))
+	})
+	// Apply remaining override flags (that didn't conflict above).
+	for of, ov := range override {
+		args = append(args, fmt.Sprintf("--%s=%s", of, ov))
+	}
+	// Add the non-flag arguments at the end.
+	args = append(args, fSet.Args()...)
 	return args
 }
 

runsc/cmd/gofer.go

@@ -159,15 +159,10 @@ func (g *Gofer) Execute(_ context.Context, f *flag.FlagSet, args ...any) subcomm
 		}
 	}
 	if g.applyCaps {
-		// Disable caps when calling myself again.
-		// Note: minimal argument handling for the default case to keep it simple.
-		args := os.Args
-		args = append(
-			args,
-			"--apply-caps=false",
-			"--setup-root=false",
-		)
-		args = append(args, g.syncFDs.flags()...)
+		overrides := g.syncFDs.flags()
+		overrides["apply-caps"] = "false"
+		overrides["setup-root"] = "false"
+		args := prepareArgs(g.Name(), f, overrides)
 		util.Fatalf("setCapsAndCallSelf(%v, %v): %v", args, goferCaps, setCapsAndCallSelf(args, goferCaps))
 		panic("unreachable")
 	}
@@ -593,11 +588,11 @@ func (g *goferSyncFDs) setFlags(f *flag.FlagSet) {
 
 // flags returns the flags necessary to pass along the current sync FD values
 // to a re-executed version of this process.
-func (g *goferSyncFDs) flags() []string {
-	return []string{
-		fmt.Sprintf("--sync-nvproxy-fd=%d", g.nvproxyFD),
-		fmt.Sprintf("--sync-userns-fd=%d", g.usernsFD),
-		fmt.Sprintf("--proc-mount-sync-fd=%d", g.procMountFD),
+func (g *goferSyncFDs) flags() map[string]string {
+	return map[string]string{
+		"sync-nvproxy-fd":    fmt.Sprintf("%d", g.nvproxyFD),
+		"sync-userns-fd":     fmt.Sprintf("%d", g.usernsFD),
+		"proc-mount-sync-fd": fmt.Sprintf("%d", g.procMountFD),
 	}
 }
 
@@ -658,20 +653,18 @@ func (g *goferSyncFDs) unmountProcfs() {
 // Postcondition: All callers must re-exec themselves after this returns,
 // unless usernsFD was -1.
 func (g *goferSyncFDs) syncUsernsForRootless() {
+	if g.usernsFD < 0 {
+		return
+	}
 	syncUsernsForRootless(g.usernsFD)
 	g.usernsFD = -1
 }
 
 // syncUsernsForRootless waits on usernsFD to be closed and then sets
 // UID/GID to 0. Note that this function calls runtime.LockOSThread().
-// This function is a no-op if usernsFD is -1.
 //
-// Postcondition: All callers must re-exec themselves after this returns,
-// unless fd is -1.
+// Postcondition: All callers must re-exec themselves after this returns.
 func syncUsernsForRootless(fd int) {
-	if fd < 0 {
-		return
-	}
 	if err := waitForFD(fd, "userns sync FD"); err != nil {
 		util.Fatalf("failed to sync on userns FD: %v", err)
 	}