Add tool to generate a Go library that executes an embedded compressed binary.

This creates a `go_template` that can generate a `go_library` rule which
extracts a compressed binary to a temporary location, then executes it.

This is useful to implement `runsc` subcommands without requiring that these
subcommands' dependencies are linked into the main `runsc` binary.

PiperOrigin-RevId: 557270682

tools/embeddedbinary/BUILD

@@ -0,0 +1,31 @@
+load("//tools:defs.bzl", "bzl_library", "go_binary")
+load("//tools/go_generics:defs.bzl", "go_template")
+
+package(
+    default_applicable_licenses = ["//:license"],
+    licenses = ["notice"],
+)
+
+bzl_library(
+    name = "defs_bzl",
+    srcs = ["defs.bzl"],
+    visibility = [
+        "//:sandbox",
+    ],
+)
+
+go_template(
+    name = "embeddedbinary_template",
+    srcs = ["embeddedbinary_template.go"],
+    visibility = [
+        "//:sandbox",
+    ],
+)
+
+go_binary(
+    name = "flatecompress",
+    srcs = ["flatecompress.go"],
+    visibility = [
+        "//:sandbox",
+    ],
+)

tools/embeddedbinary/README.md

@@ -0,0 +1,30 @@
+# embeddedbinary
+
+`embeddedbinary` can embed a binary inside a Go binary, and provides functions
+to execute it.
+
+Embedded binaries are compressed to save on size. They require temporary disk
+space to execute, but the disk space is automatically freed when the child
+program exits.
+
+## Usage
+
+Use the `embedded_binary_go_library` rule defined in `defs.bzl`.
+
+```build
+load(".../defs.bzl", "embedded_binary_go_library")
+
+# Declare a binary target:
+go_binary(
+    name = "my_binary",
+    srcs = ["my_binary.go"],
+)
+
+# Generate a go_library rule that can execute the binary target:
+embedded_binary_go_library(
+    name = "my_library",
+    binary = ":my_binary",
+)
+```
+
+See `test/BUILD` under this directory for a full example.

tools/embeddedbinary/defs.bzl

@@ -0,0 +1,60 @@
+"""BUILD rule for embedded binaries."""
+
+load("//tools:defs.bzl", "go_library")
+load("//tools/go_generics:defs.bzl", "go_template_instance")
+
+_EMBEDDED_BINARY_TEMPLATE = "//tools/embeddedbinary:embeddedbinary_template"
+_FLATECOMPRESS = "//tools/embeddedbinary:flatecompress"
+
+def embedded_binary_go_library(
+        name,
+        binary,
+        binary_name = None,
+        out = None,
+        go_package_name = None,
+        visibility = None):
+    """Embed a binary and generate a go_library target that can execute it.
+
+    The binary will be compressed, and needs temporary space to be available
+    when executing.
+
+    Args:
+        name: The name of the go_library rule.
+        binary: Binary BUILD target that should be embedded.
+        binary_name: The name (i.e. typical argv[0]) of the binary being
+          embedded, defaults to `name`.
+        out: Output filename of the go_library rule, defaults to `name + ".go"`.
+        go_package_name: Package name of the go_library, defaults to `name`.
+        visibility: Visibility of the go_library rule.
+    """
+    if binary_name == None:
+        binary_name = name
+    if out == None:
+        out = name + ".go"
+    if go_package_name == None:
+        go_package_name = name
+    compressed_binary = binary_name + ".flate"
+    native.genrule(
+        name = name + "_flate",
+        outs = [compressed_binary],
+        cmd = "$(location %s) < $(SRCS) > $(OUTS)" % (_FLATECOMPRESS,),
+        srcs = [binary],
+        tools = [_FLATECOMPRESS],
+    )
+    go_template_instance(
+        name = name + "_lib",
+        template = _EMBEDDED_BINARY_TEMPLATE,
+        package = go_package_name,
+        out = out,
+        substrs = {
+            "embedded.bin.name": binary_name,
+            "//go:embed embedded.bin.flate": "//go:embed %s" % (compressed_binary,),
+        },
+    )
+    go_library(
+        name = name,
+        srcs = [out],
+        embedsrcs = [compressed_binary],
+        deps = ["@org_golang_x_sys//unix:go_default_library"],
+        visibility = visibility,
+    )

tools/embeddedbinary/embeddedbinary_template.go

@@ -0,0 +1,129 @@
+// Copyright 2023 The gVisor Authors.
+//
+// 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 embeddedbinary embeds an external binary and provides a function to
+// exec it.
+package embeddedbinary
+
+import (
+	"bytes"
+	"compress/flate"
+	_ "embed"
+	"fmt"
+	"io"
+	"os"
+	"path"
+	"runtime"
+	"syscall"
+
+	"golang.org/x/sys/unix"
+)
+
+// BinaryName is the name of the embedded binary.
+const BinaryName = "embedded.bin.name"
+
+//go:embed embedded.bin.flate
+var compressedBinary []byte
+
+// Options is the set of options to execute the embedded binary.
+type Options struct {
+	// Argv is the set of arguments to exec with.
+	// `Argv[0]` is the name of the binary as invoked.
+	// If Argv is empty, it will default to a single-element slice, with
+	// `Argv[0]` being the binary name.
+	Argv []string
+
+	// Envv is the set of environment variables to pass to the executed process.
+	Envv []string
+
+	// Files is the set of file descriptors to pass to forked processes.
+	// Only used when forking, not pure exec'ing.
+	Files []uintptr
+}
+
+// run decompresses and run the embedded binary with the given arguments.
+// If fork is true, the binary runs in a separate process, and its PID is
+// returned.
+// Otherwise, the binary is exec'd, so the current process stops executing.
+func run(options Options, fork bool) (int, error) {
+	if len(options.Argv) == 0 {
+		options.Argv = []string{BinaryName}
+	}
+	decompressed := flate.NewReader(bytes.NewReader(compressedBinary))
+	runtime.LockOSThread()
+	defer runtime.UnlockOSThread()
+	myPID := os.Getpid()
+	oldMask := unix.Umask(0077)
+	defer unix.Umask(oldMask)
+	tmpDir, err := os.MkdirTemp("", "gvisor.*.tmp")
+	if err != nil {
+		return 0, fmt.Errorf("cannot create temp directory: %w", err)
+	}
+	tmpDirHandle, err := os.Open(tmpDir)
+	if err != nil {
+		return 0, fmt.Errorf("cannot open temp directory: %w", err)
+	}
+	defer tmpDirHandle.Close()
+	binPath := path.Join(tmpDir, BinaryName)
+	tmpFile, err := os.OpenFile(binPath, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0700)
+	if err != nil {
+		return 0, fmt.Errorf("cannot open temp file: %w", err)
+	}
+	if err := os.RemoveAll(tmpDir); err != nil {
+		return 0, fmt.Errorf("cannot remove temp directory: %w", err)
+	}
+	unix.Umask(oldMask)
+	if _, err := io.Copy(tmpFile, decompressed); err != nil {
+		tmpFile.Close()
+		return 0, fmt.Errorf("cannot decompress embedded binary or write it to temporary file: %w", err)
+	}
+	// Reopen the file for reading.
+	tmpFileReadOnly, err := os.OpenFile(fmt.Sprintf("/proc/%d/fd/%d", myPID, tmpFile.Fd()), os.O_RDONLY, 0700)
+	if err != nil {
+		tmpFile.Close()
+		return 0, fmt.Errorf("cannot re-open temp file for reading: %w", err)
+	}
+	if err := tmpFile.Close(); err != nil {
+		return 0, fmt.Errorf("cannot close temp file: %w", err)
+	}
+	defer tmpFileReadOnly.Close()
+	tmpFD := tmpFileReadOnly.Fd()
+	if _, err := unix.Seek(int(tmpFD), 0, unix.SEEK_SET); err != nil {
+		return 0, fmt.Errorf("cannot seek temp file back to 0: %w", err)
+	}
+	fdPath := fmt.Sprintf("/proc/%d/fd/%d", myPID, tmpFD)
+	if fork {
+		return syscall.ForkExec(fdPath, options.Argv, &syscall.ProcAttr{
+			Env:   options.Envv,
+			Files: options.Files,
+		})
+	}
+	if err := unix.Exec(fdPath, options.Argv, options.Envv); err != nil {
+		return 0, fmt.Errorf("cannot exec embedded binary: %w", err)
+	}
+	panic("unreachable")
+}
+
+// Exec execs the embedded binary. The current process is replaced.
+// This function only returns if unsuccessful.
+func Exec(options Options) error {
+	_, err := run(options, false)
+	return err
+}
+
+// ForkExec runs the embedded binary in a separate process.
+// Returns the PID of the child process.
+func ForkExec(options Options) (int, error) {
+	return run(options, true)
+}

tools/embeddedbinary/flatecompress.go

@@ -0,0 +1,44 @@
+// Copyright 2023 The gVisor Authors.
+//
+// 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.
+
+// flatecompress compresses data from stdin and writes it to stdout with flate.
+package main
+
+import (
+	"compress/flate"
+	"fmt"
+	"io"
+	"os"
+)
+
+func main() {
+	writer, err := flate.NewWriter(os.Stdout, flate.BestCompression)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Cannot create flate writer: %v\n", err)
+		os.Exit(1)
+	}
+	if _, err := io.Copy(writer, os.Stdin); err != nil {
+		fmt.Fprintf(os.Stderr, "Failed to compress binary: %v\n", err)
+		os.Exit(1)
+	}
+	if err := writer.Flush(); err != nil {
+		fmt.Fprintf(os.Stderr, "Cannot flush: %v\n", err)
+		os.Exit(1)
+	}
+	if err := writer.Close(); err != nil {
+		fmt.Fprintf(os.Stderr, "Cannot close writer: %v\n", err)
+		os.Exit(1)
+	}
+	os.Exit(0)
+}

tools/embeddedbinary/test/BUILD

@@ -0,0 +1,27 @@
+load("//tools:defs.bzl", "go_binary", "go_test")
+
+package(
+    default_applicable_licenses = ["//:license"],
+    licenses = ["notice"],
+)
+
+# helloworld_bundler is a Go program that imports helloworld/helloworld.go and
+# calls its functions. Therefore, it should print "Hello, gVisor!\n" to stdout,
+# as helloworld/helloworld_bundlee.go does.
+go_binary(
+    name = "helloworld_bundler",
+    srcs = ["helloworld_bundler.go"],
+    deps = [
+        "//tools/embeddedbinary/test/helloworld",
+        "@org_golang_x_sys//unix:go_default_library",
+    ],
+)
+
+# helloworld_test is a test that executes helloworld_bundler as a subprocess.
+# It verifies that its output is "Hello, gVisor!\n".
+go_test(
+    name = "helloworld_test",
+    srcs = ["helloworld_test.go"],
+    data = [":helloworld_bundler"],
+    deps = ["//pkg/test/testutil"],
+)

tools/embeddedbinary/test/helloworld/BUILD

@@ -0,0 +1,24 @@
+load("//tools:defs.bzl", "go_binary")
+load("//tools/embeddedbinary:defs.bzl", "embedded_binary_go_library")
+
+package(
+    default_applicable_licenses = ["//:license"],
+    licenses = ["notice"],
+)
+
+# helloworld_bundlee is a simple Go program that prints "Hello, gVisor!\n"
+# to stdout. It is bundled by other rules.
+go_binary(
+    name = "helloworld_bundlee",
+    srcs = ["helloworld_bundlee.go"],
+)
+
+# helloworld generates a Go source file called "helloworld.go" which embeds
+# the helloworld_bundlee program.
+embedded_binary_go_library(
+    name = "helloworld",
+    binary = ":helloworld_bundlee",
+    visibility = [
+        "//tools/embeddedbinary/test:__subpackages__",
+    ],
+)

tools/embeddedbinary/test/helloworld/helloworld_bundlee.go

@@ -0,0 +1,27 @@
+// Copyright 2023 The gVisor Authors.
+//
+// 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.
+
+// helloworld_bundlee writes "Hello, gVisor!\n" to stdout.
+// It is meant to be bundled into helloworld_bundler.
+package main
+
+import (
+	"fmt"
+	"os"
+)
+
+func main() {
+	fmt.Fprintf(os.Stdout, "Hello, gVisor!\n")
+	os.Exit(0)
+}

tools/embeddedbinary/test/helloworld_bundler.go

@@ -0,0 +1,67 @@
+// Copyright 2023 The gVisor Authors.
+//
+// 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.
+
+// helloworld_bundler bundles helloworld_bundlee and executes it.
+package main
+
+import (
+	"fmt"
+	"os"
+
+	syscall "golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/tools/embeddedbinary/test/helloworld"
+)
+
+func doExec() {
+	if err := helloworld.Exec(helloworld.Options{}); err != nil {
+		fmt.Fprintf(os.Stderr, "Failed to exec embedded binary: %v\n", err)
+		os.Exit(1)
+	}
+	fmt.Fprintf(os.Stderr, "Unreachable\n")
+	os.Exit(1)
+}
+
+func doForkExec() {
+	childPID, err := helloworld.ForkExec(helloworld.Options{
+		// Share stdin/stdout/stderr with child process.
+		Files: []uintptr{os.Stdin.Fd(), os.Stdout.Fd(), os.Stderr.Fd()},
+	})
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Failed to fork+exec embedded binary: %v\n", err)
+		os.Exit(1)
+	}
+	var waitStatus syscall.WaitStatus
+	if _, err := syscall.Wait4(childPID, &waitStatus, 0, nil); err != nil {
+		fmt.Fprintf(os.Stderr, "Failed to wait for child embedded binary: %v\n", err)
+		os.Exit(1)
+	}
+	if status := waitStatus.ExitStatus(); status != 0 {
+		fmt.Fprintf(os.Stderr, "Child embedded binary returned code %d\n", status)
+		os.Exit(status)
+	}
+	os.Exit(0)
+}
+
+func main() {
+	for _, arg := range os.Args {
+		switch arg {
+		case "--mode=exec":
+			doExec()
+		case "--mode=fork":
+			doForkExec()
+		}
+	}
+	fmt.Fprintf(os.Stderr, "Must specify either --mode=exec or --mode=fork.\n")
+	os.Exit(1)
+}

tools/embeddedbinary/test/helloworld_test.go

@@ -0,0 +1,49 @@
+// Copyright 2023 The gVisor Authors.
+//
+// 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 helloworld_test
+
+import (
+	"context"
+	"fmt"
+	"os/exec"
+	"testing"
+	"time"
+
+	"gvisor.dev/gvisor/pkg/test/testutil"
+)
+
+// TestHelloworld executes helloworld_bundler and verifies that its output
+// matches "Hello, gVisor!\n".
+func TestHelloworld(t *testing.T) {
+	ctx, ctxCancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer ctxCancel()
+	helloWorldPath, err := testutil.FindFile("tools/embeddedbinary/test/helloworld_bundler")
+	if err != nil {
+		t.Fatalf("Cannot find helloworld_bundler path: %v", err)
+	}
+	for _, mode := range []string{"exec", "fork"} {
+		t.Run(mode, func(t *testing.T) {
+			output, err := exec.CommandContext(ctx, helloWorldPath, fmt.Sprintf("--mode=%s", mode)).CombinedOutput()
+			outputStr := string(output)
+			if err != nil {
+				t.Fatalf("Failed to execute helloworld_bundler: %v; output:\n%v\n", err, outputStr)
+			}
+			want := "Hello, gVisor!\n"
+			if outputStr != want {
+				t.Fatalf("helloworld_bundler: got output %q want %q", outputStr, want)
+			}
+		})
+	}
+}