From 6d021e7eb2112e4692470c3a60a9a028adbe7e37 Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Tue, 18 Jun 2024 13:26:46 -0700
Subject: [PATCH] Bundle `libioctl_hook.so` into the `run_sniffer` tool.

This makes it more convenient to use the `run_sniffer` tool as it is now
self-contained.

This also removes the use of RELR relocations for `libioctl_hook.so` as
I've found that they don't work on the `nvidia-smi` binary I'm working with.
This improves compatibility in general.

PiperOrigin-RevId: 644493298
---
 tools/ioctl_sniffer/BUILD          | 13 ++++++++++
 tools/ioctl_sniffer/README.md      |  5 ----
 tools/ioctl_sniffer/run_sniffer.go | 38 ++++++++++++++++++++++++++++--
 3 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/tools/ioctl_sniffer/BUILD b/tools/ioctl_sniffer/BUILD
index 3d4baef98..7b06a640c 100644
--- a/tools/ioctl_sniffer/BUILD
+++ b/tools/ioctl_sniffer/BUILD
@@ -19,6 +19,10 @@ cc_binary(
         "sniffer_bridge.cc",
         "sniffer_bridge.h",
     ],
+    # Needed to support being used when LD_PRELOAD'd into binaries that are
+    # built with a libc library that doesn't support RELR relocations
+    # (such as nvidia-smi).
+    features = ["-enable_relr"],
     linkshared = True,
     deps = [
         ":ioctl_cc_proto",
@@ -31,6 +35,15 @@ cc_binary(
 go_binary(
     name = "run_sniffer",
     srcs = ["run_sniffer.go"],
+    embedsrcs = [
+        # The 'keep' comment is needed to prevent glaze from removing this
+        # dependency. This is because the `:ioctl_hook` `cc_binary` rule
+        # output file is 'libioctl_hook.so' (because of `linkshared = True`),
+        # but glaze assumes that the output file is just `ioctl_hook`, so it
+        # can't associate the `go:embed` statement in `run_sniffer.go` with
+        # `:ioctl_hook`'s actual output file.
+        ":ioctl_hook",  # keep
+    ],
     static = True,
     deps = [
         "//pkg/log",
diff --git a/tools/ioctl_sniffer/README.md b/tools/ioctl_sniffer/README.md
index f1ab9a2da..edc9859a2 100644
--- a/tools/ioctl_sniffer/README.md
+++ b/tools/ioctl_sniffer/README.md
@@ -19,7 +19,6 @@ To start, we need to build the shared library and Go binary:
 
 ```
 make copy TARGETS=//tools/ioctl_sniffer:run_sniffer DESTINATION=bin/
-make copy TARGETS=//tools/ioctl_sniffer:ioctl_hook DESTINATION=bin/
 ```
 
 Once we have the binary, we can hook into any GPU workload by passing the
@@ -46,7 +45,3 @@ Alloc:
     ...
 Unknown:
 ```
-
-Note that by default, `run_sniffer` assumes the shared library is located in the
-same directory. You can specify the path to the library with the optional
-`-ld_preload` flag.
diff --git a/tools/ioctl_sniffer/run_sniffer.go b/tools/ioctl_sniffer/run_sniffer.go
index 6af7bbb8a..1c91e2412 100644
--- a/tools/ioctl_sniffer/run_sniffer.go
+++ b/tools/ioctl_sniffer/run_sniffer.go
@@ -23,9 +23,31 @@ import (
 
 	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/tools/ioctl_sniffer/sniffer"
+
+	_ "embed" // Necessary to use go:embed.
 )
 
-var ldPreloadPath = flag.String("ld_preload", "./libioctl_hook.so", "The path to the LD_PRELOAD library.")
+//go:embed libioctl_hook.so
+var ioctlHookSharedObject []byte
+
+// createSharedObject creates a temporary directory containing the ioctl hook
+// shared object, and returns the path to it. This file will be automatically
+// deleted when the program exits.
+func createSharedObject() (*os.File, error) {
+	tmpFile, err := os.CreateTemp(os.TempDir(), "libioctl_hook.*.so")
+	if err != nil {
+		return nil, fmt.Errorf("failed to create temporary file: %w", err)
+	}
+	// Remove the file from the filesystem but keep a handle open to it
+	// so that it lasts only as long as the program does.
+	if err := os.Remove(tmpFile.Name()); err != nil {
+		return nil, fmt.Errorf("failed to unlink temporary file: %w", err)
+	}
+	if _, err := tmpFile.Write(ioctlHookSharedObject); err != nil {
+		return nil, fmt.Errorf("failed to write to temporary file: %w", err)
+	}
+	return tmpFile, nil
+}
 
 // Main is our main function.
 func Main() error {
@@ -39,6 +61,16 @@ func Main() error {
 		return fmt.Errorf("failed to init sniffer: %w", err)
 	}
 
+	hookFile, err := createSharedObject()
+	if err != nil {
+		return fmt.Errorf("failed to create shared object file: %w", err)
+	}
+	defer func() {
+		if err := hookFile.Close(); err != nil {
+			log.Warningf("failed to close shared object file: %w", err)
+		}
+	}()
+
 	// Create a pipe to read the output of the command.
 	r, w, err := os.Pipe()
 	if err != nil {
@@ -51,7 +83,9 @@ func Main() error {
 	cmd.Stdin = os.Stdin
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
-	cmd.Env = append(os.Environ(), fmt.Sprintf("LD_PRELOAD=%s", *ldPreloadPath))
+	// Refer to the hook file by file descriptor here as its named file no
+	// longer exists.
+	cmd.Env = append(os.Environ(), fmt.Sprintf("LD_PRELOAD=/proc/%d/fd/%d", os.Getpid(), hookFile.Fd()))
 
 	// Run the command and start reading the output.
 	if err := cmd.Start(); err != nil {