Add nsjail support to genrule

By setting `use_nsjail: true`, nsjail will be used instead of soong
sbox. The usage of nsjail must be highly restricted for now.

Bug: 364178791
Test: m lk.bin
Change-Id: Iffd1404093a1165b4384fdec58525fe16f7ed7af

1 files changed, 167 insertions, 19 deletions

diff --git a/android/rule_builder.go b/android/rule_builder.go
index 18bbcab5c..56de9cd00 100644
--- a/android/rule_builder.go
+++ b/android/rule_builder.go
@@ -38,6 +38,9 @@ const sboxOutSubDir = "out"
 const sboxToolsSubDir = "tools"
 const sboxOutDir = sboxSandboxBaseDir + "/" + sboxOutSubDir
 
+const nsjailToolsSubDir = "tools"
+const nsjailOutDir = "out"
+
 // RuleBuilder provides an alternative to ModuleContext.Rule and ModuleContext.Build to add a command line to the build
 // graph.
 type RuleBuilder struct {
@@ -59,6 +62,9 @@ type RuleBuilder struct {
 	sboxManifestPath WritablePath
 	missingDeps      []string
 	args             map[string]string
+	nsjail           bool
+	nsjailBasePath   WritablePath
+	nsjailImplicits  Paths
 }
 
 // NewRuleBuilder returns a newly created RuleBuilder.
@@ -165,12 +171,43 @@ func (r *RuleBuilder) Sbox(outputDir WritablePath, manifestPath WritablePath) *R
 	if len(r.commands) > 0 {
 		panic("Sbox() may not be called after Command()")
 	}
+	if r.nsjail {
+		panic("Sbox() may not be called after Nsjail()")
+	}
 	r.sbox = true
 	r.outDir = outputDir
 	r.sboxManifestPath = manifestPath
 	return r
 }
 
+// Nsjail marks the rule as needing to be wrapped by nsjail. The outputDir should point to the
+// output directory that nsjail will mount to out/. It should not be written to by any other rule.
+// baseDir should point to a location where nsjail will mount to /nsjail_build_sandbox, which will
+// be the working directory of the command.
+func (r *RuleBuilder) Nsjail(outputDir WritablePath, baseDir WritablePath) *RuleBuilder {
+	if len(r.commands) > 0 {
+		panic("Nsjail() may not be called after Command()")
+	}
+	if r.sbox {
+		panic("Nsjail() may not be called after Sbox()")
+	}
+	r.nsjail = true
+	r.outDir = outputDir
+	r.nsjailBasePath = baseDir
+	return r
+}
+
+// NsjailImplicits adds implicit inputs that are not directly mounted. This is useful when
+// the rule mounts directories, as files within those directories can be globbed and
+// tracked as dependencies with NsjailImplicits().
+func (r *RuleBuilder) NsjailImplicits(inputs Paths) *RuleBuilder {
+	if !r.nsjail {
+		panic("NsjailImplicits() must be called after Nsjail()")
+	}
+	r.nsjailImplicits = append(r.nsjailImplicits, inputs...)
+	return r
+}
+
 // SandboxTools enables tool sandboxing for the rule by copying any referenced tools into the
 // sandbox.
 func (r *RuleBuilder) SandboxTools() *RuleBuilder {
@@ -514,7 +551,73 @@ func (r *RuleBuilder) build(name string, desc string, ninjaEscapeCommandString b
 
 	commandString := strings.Join(commands, " && ")
 
-	if r.sbox {
+	if !r.sbox {
+		// If not using sbox the rule will run the command directly, put the hash of the
+		// list of input files in a comment at the end of the command line to ensure ninja
+		// reruns the rule when the list of input files changes.
+		commandString += " # hash of input list: " + hashSrcFiles(inputs)
+	}
+
+	if r.nsjail {
+		var nsjailCmd strings.Builder
+		nsjailPath := r.ctx.Config().PrebuiltBuildTool(r.ctx, "nsjail")
+		nsjailCmd.WriteString("mkdir -p ")
+		nsjailCmd.WriteString(r.nsjailBasePath.String())
+		nsjailCmd.WriteString(" && ")
+		nsjailCmd.WriteString(nsjailPath.String())
+		nsjailCmd.WriteRune(' ')
+		nsjailCmd.WriteString("-B $PWD/")
+		nsjailCmd.WriteString(r.nsjailBasePath.String())
+		nsjailCmd.WriteString(":nsjail_build_sandbox")
+
+		// out is mounted to $(genDir).
+		nsjailCmd.WriteString(" -B $PWD/")
+		nsjailCmd.WriteString(r.outDir.String())
+		nsjailCmd.WriteString(":nsjail_build_sandbox/out")
+
+		for _, input := range inputs {
+			nsjailCmd.WriteString(" -R $PWD/")
+			nsjailCmd.WriteString(input.String())
+			nsjailCmd.WriteString(":nsjail_build_sandbox/")
+			nsjailCmd.WriteString(r.nsjailPathForInputRel(input))
+		}
+		for _, tool := range tools {
+			nsjailCmd.WriteString(" -R $PWD/")
+			nsjailCmd.WriteString(tool.String())
+			nsjailCmd.WriteString(":nsjail_build_sandbox/")
+			nsjailCmd.WriteString(nsjailPathForToolRel(r.ctx, tool))
+		}
+		inputs = append(inputs, tools...)
+		for _, c := range r.commands {
+			for _, tool := range c.packagedTools {
+				nsjailCmd.WriteString(" -R $PWD/")
+				nsjailCmd.WriteString(tool.srcPath.String())
+				nsjailCmd.WriteString(":nsjail_build_sandbox/")
+				nsjailCmd.WriteString(nsjailPathForPackagedToolRel(tool))
+				inputs = append(inputs, tool.srcPath)
+			}
+		}
+
+		// These five directories are necessary to run native host tools like /bin/bash and py3-cmd.
+		nsjailCmd.WriteString(" -R /bin")
+		nsjailCmd.WriteString(" -R /lib")
+		nsjailCmd.WriteString(" -R /lib64")
+		nsjailCmd.WriteString(" -R /dev")
+		nsjailCmd.WriteString(" -R /usr")
+
+		nsjailCmd.WriteString(" -m none:/tmp:tmpfs:size=1073741824") // 1GB, should be enough
+		nsjailCmd.WriteString(" -D nsjail_build_sandbox")
+		nsjailCmd.WriteString(" --disable_rlimits")
+		nsjailCmd.WriteString(" -q")
+		nsjailCmd.WriteString(" -- ")
+		nsjailCmd.WriteString("/bin/bash -c ")
+		nsjailCmd.WriteString(proptools.ShellEscape(commandString))
+
+		commandString = nsjailCmd.String()
+
+		inputs = append(inputs, nsjailPath)
+		inputs = append(inputs, r.nsjailImplicits...)
+	} else if r.sbox {
 		// If running the command inside sbox, write the rule data out to an sbox
 		// manifest.textproto.
 		manifest := sbox_proto.Manifest{}
@@ -734,11 +837,6 @@ func (r *RuleBuilder) build(name string, desc string, ninjaEscapeCommandString b
 			rewrapperCommand := r.rbeParams.NoVarTemplate(r.ctx.Config().RBEWrapper())
 			commandString = rewrapperCommand + " bash -c '" + strings.ReplaceAll(commandString, `'`, `'\''`) + "'"
 		}
-	} else {
-		// If not using sbox the rule will run the command directly, put the hash of the
-		// list of input files in a comment at the end of the command line to ensure ninja
-		// reruns the rule when the list of input files changes.
-		commandString += " # hash of input list: " + hashSrcFiles(inputs)
 	}
 
 	// Ninja doesn't like multiple outputs when depfiles are enabled, move all but the first output to
@@ -869,6 +967,8 @@ func (c *RuleBuilderCommand) PathForInput(path Path) string {
 			rel = filepath.Join(sboxSandboxBaseDir, rel)
 		}
 		return rel
+	} else if c.rule.nsjail {
+		return c.rule.nsjailPathForInputRel(path)
 	}
 	return path.String()
 }
@@ -894,6 +994,10 @@ func (c *RuleBuilderCommand) PathForOutput(path WritablePath) string {
 		// Errors will be handled in RuleBuilder.Build where we have a context to report them
 		rel, _, _ := maybeRelErr(c.rule.outDir.String(), path.String())
 		return filepath.Join(sboxOutDir, rel)
+	} else if c.rule.nsjail {
+		// Errors will be handled in RuleBuilder.Build where we have a context to report them
+		rel, _, _ := maybeRelErr(c.rule.outDir.String(), path.String())
+		return filepath.Join(nsjailOutDir, rel)
 	}
 	return path.String()
 }
@@ -945,15 +1049,49 @@ func sboxPathForPackagedToolRel(spec PackagingSpec) string {
 	return filepath.Join(sboxToolsSubDir, "out", spec.relPathInPackage)
 }
 
+func nsjailPathForToolRel(ctx BuilderContext, path Path) string {
+	// Errors will be handled in RuleBuilder.Build where we have a context to report them
+	toolDir := pathForInstall(ctx, ctx.Config().BuildOS, ctx.Config().BuildArch, "")
+	relOutSoong, isRelOutSoong, _ := maybeRelErr(toolDir.String(), path.String())
+	if isRelOutSoong {
+		// The tool is in the Soong output directory, it will be copied to __SBOX_OUT_DIR__/tools/out
+		return filepath.Join(nsjailToolsSubDir, "out", relOutSoong)
+	}
+	// The tool is in the source directory, it will be copied to __SBOX_OUT_DIR__/tools/src
+	return filepath.Join(nsjailToolsSubDir, "src", path.String())
+}
+
+func (r *RuleBuilder) nsjailPathForInputRel(path Path) string {
+	rel, isRelSboxOut, _ := maybeRelErr(r.outDir.String(), path.String())
+	if isRelSboxOut {
+		return filepath.Join(nsjailOutDir, rel)
+	}
+	return path.String()
+}
+
+func (r *RuleBuilder) nsjailPathsForInputsRel(paths Paths) []string {
+	ret := make([]string, len(paths))
+	for i, path := range paths {
+		ret[i] = r.nsjailPathForInputRel(path)
+	}
+	return ret
+}
+
+func nsjailPathForPackagedToolRel(spec PackagingSpec) string {
+	return filepath.Join(nsjailToolsSubDir, "out", spec.relPathInPackage)
+}
+
 // PathForPackagedTool takes a PackageSpec for a tool and returns the corresponding path for the
 // tool after copying it into the sandbox.  This can be used  on the RuleBuilder command line to
 // reference the tool.
 func (c *RuleBuilderCommand) PathForPackagedTool(spec PackagingSpec) string {
-	if !c.rule.sboxTools {
-		panic("PathForPackagedTool() requires SandboxTools()")
+	if c.rule.sboxTools {
+		return filepath.Join(sboxSandboxBaseDir, sboxPathForPackagedToolRel(spec))
+	} else if c.rule.nsjail {
+		return nsjailPathForPackagedToolRel(spec)
+	} else {
+		panic("PathForPackagedTool() requires SandboxTools() or Nsjail()")
 	}
-
-	return filepath.Join(sboxSandboxBaseDir, sboxPathForPackagedToolRel(spec))
 }
 
 // PathForTool takes a path to a tool, which may be an output file or a source file, and returns
@@ -962,6 +1100,8 @@ func (c *RuleBuilderCommand) PathForPackagedTool(spec PackagingSpec) string {
 func (c *RuleBuilderCommand) PathForTool(path Path) string {
 	if c.rule.sbox && c.rule.sboxTools {
 		return filepath.Join(sboxSandboxBaseDir, sboxPathForToolRel(c.rule.ctx, path))
+	} else if c.rule.nsjail {
+		return nsjailPathForToolRel(c.rule.ctx, path)
 	}
 	return path.String()
 }
@@ -976,6 +1116,12 @@ func (c *RuleBuilderCommand) PathsForTools(paths Paths) []string {
 			ret = append(ret, filepath.Join(sboxSandboxBaseDir, sboxPathForToolRel(c.rule.ctx, path)))
 		}
 		return ret
+	} else if c.rule.nsjail {
+		var ret []string
+		for _, path := range paths {
+			ret = append(ret, nsjailPathForToolRel(c.rule.ctx, path))
+		}
+		return ret
 	}
 	return paths.Strings()
 }
@@ -983,20 +1129,22 @@ func (c *RuleBuilderCommand) PathsForTools(paths Paths) []string {
 // PackagedTool adds the specified tool path to the command line.  It can only be used with tool
 // sandboxing enabled by SandboxTools(), and will copy the tool into the sandbox.
 func (c *RuleBuilderCommand) PackagedTool(spec PackagingSpec) *RuleBuilderCommand {
-	if !c.rule.sboxTools {
-		panic("PackagedTool() requires SandboxTools()")
-	}
-
 	c.packagedTools = append(c.packagedTools, spec)
-	c.Text(sboxPathForPackagedToolRel(spec))
+	if c.rule.sboxTools {
+		c.Text(sboxPathForPackagedToolRel(spec))
+	} else if c.rule.nsjail {
+		c.Text(nsjailPathForPackagedToolRel(spec))
+	} else {
+		panic("PackagedTool() requires SandboxTools() or Nsjail()")
+	}
 	return c
 }
 
 // ImplicitPackagedTool copies the specified tool into the sandbox without modifying the command
 // line.  It can only be used with tool sandboxing enabled by SandboxTools().
 func (c *RuleBuilderCommand) ImplicitPackagedTool(spec PackagingSpec) *RuleBuilderCommand {
-	if !c.rule.sboxTools {
-		panic("ImplicitPackagedTool() requires SandboxTools()")
+	if !c.rule.sboxTools && !c.rule.nsjail {
+		panic("ImplicitPackagedTool() requires SandboxTools() or Nsjail()")
 	}
 
 	c.packagedTools = append(c.packagedTools, spec)
@@ -1006,8 +1154,8 @@ func (c *RuleBuilderCommand) ImplicitPackagedTool(spec PackagingSpec) *RuleBuild
 // ImplicitPackagedTools copies the specified tools into the sandbox without modifying the command
 // line.  It can only be used with tool sandboxing enabled by SandboxTools().
 func (c *RuleBuilderCommand) ImplicitPackagedTools(specs []PackagingSpec) *RuleBuilderCommand {
-	if !c.rule.sboxTools {
-		panic("ImplicitPackagedTools() requires SandboxTools()")
+	if !c.rule.sboxTools && !c.rule.nsjail {
+		panic("ImplicitPackagedTools() requires SandboxTools() or Nsjail()")
 	}
 
 	c.packagedTools = append(c.packagedTools, specs...)