Merge pull request #60 from ilmanzo/improve_wrap_unwrap

Allow wrap/unwrap many binaries at once - give useful error on finding the .so library

Author: ilmanzo

cmd/funkoverage.go

@@ -8,7 +8,7 @@ import (
 	"strings"
 )
 
-const versionString = "0.4.2"
+const versionString = "0.4.3"
 
 // --- CLI ---
 
@@ -34,20 +34,20 @@ func main() {
 	case "wrap", "-w":
 		wrapCmd.Parse(os.Args[2:])
 		if wrapCmd.NArg() < 1 {
-			fmt.Println("wrap: missing binary path")
+			fmt.Println("wrap: missing binary path(s)")
 			os.Exit(1)
 		}
-		if err := wrap(wrapCmd.Arg(0)); err != nil {
+		if err := wrapMany(wrapCmd.Args()); err != nil {
 			fmt.Println("wrap error:", err)
 			os.Exit(1)
 		}
 	case "unwrap", "-u":
 		unwrapCmd.Parse(os.Args[2:])
 		if unwrapCmd.NArg() < 1 {
-			fmt.Println("unwrap: missing binary path")
+			fmt.Println("unwrap: missing binary path(s)")
 			os.Exit(1)
 		}
-		if err := unwrap(unwrapCmd.Arg(0)); err != nil {
+		if err := unwrapMany(unwrapCmd.Args()); err != nil {
 			fmt.Println("unwrap error:", err)
 			os.Exit(1)
 		}

cmd/funkoverage_test.go

@@ -173,6 +173,57 @@ func TestWrapUnwrapLogic(t *testing.T) {
 	}
 }
 
+func TestWrapManyAndUnwrapMany(t *testing.T) {
+	tmp := t.TempDir()
+	os.Setenv("PIN_ROOT", "/tmp/pin")
+	os.Setenv("PIN_TOOL_SEARCH_DIR", tmp)
+	os.Setenv("SAFE_BIN_DIR", tmp)
+	os.Setenv("LOG_DIR", tmp)
+	// Create dummy FuncTracer.so
+	funcTracer := filepath.Join(tmp, "FuncTracer.so")
+	if err := os.WriteFile(funcTracer, []byte("dummy"), 0644); err != nil {
+		t.Fatal(err)
+	}
+
+	// Create multiple fake ELF binaries
+	bin1 := filepath.Join(tmp, "bin1")
+	bin2 := filepath.Join(tmp, "bin2")
+	bin3 := filepath.Join(tmp, "bin3")
+	for _, bin := range []string{bin1, bin2, bin3} {
+		if err := os.WriteFile(bin, []byte("\x7fELFfoobar"), 0755); err != nil {
+			t.Fatalf("failed to create %s: %v", bin, err)
+		}
+	}
+
+	// Wrap all binaries
+	if err := wrapMany([]string{bin1, bin2, bin3}); err != nil {
+		t.Fatalf("wrapMany failed: %v", err)
+	}
+	for _, bin := range []string{bin1, bin2, bin3} {
+		content, err := os.ReadFile(bin)
+		if err != nil {
+			t.Fatalf("failed to read wrapped binary %s: %v", bin, err)
+		}
+		if !strings.Contains(string(content), wrapperIDComment) {
+			t.Errorf("binary %s was not wrapped", bin)
+		}
+	}
+
+	// Unwrap all binaries
+	if err := unwrapMany([]string{bin1, bin2, bin3}); err != nil {
+		t.Fatalf("unwrapMany failed: %v", err)
+	}
+	for _, bin := range []string{bin1, bin2, bin3} {
+		_, err := os.ReadFile(bin)
+		if err != nil {
+			t.Fatalf("failed to read unwrapped binary %s: %v", bin, err)
+		}
+		if !isELF(bin) {
+			t.Errorf("binary %s was not restored to ELF", bin)
+		}
+	}
+}
+
 func TestGenerateHTMLReportBaseName(t *testing.T) {
 	tmp := t.TempDir()
 	data := &CoverageData{

cmd/wrapunwrap.go

@@ -10,7 +10,12 @@ import (
 	"time"
 )
 
-const wrapperIDComment = "# Pin Wrapper generated by Go tool funkoverage"
+const (
+	wrapperIDComment        = "# Pin Wrapper generated by Go tool funkoverage"
+	defaultPinToolSearchDir = "/usr/lib64/coverage-tools"
+	defaultLogDir           = "/var/coverage/data"
+	defaultSafeBinDir       = "/var/coverage/bin"
+)
 
 // --- Wrapper Management ---
 
@@ -73,15 +78,15 @@ func wrap(targetBinary string) error {
 	}
 	PIN_TOOL_SEARCH_DIR := os.Getenv("PIN_TOOL_SEARCH_DIR")
 	if PIN_TOOL_SEARCH_DIR == "" {
-		PIN_TOOL_SEARCH_DIR = "/usr/lib64/coverage-tools"
+		PIN_TOOL_SEARCH_DIR = defaultPinToolSearchDir
 	}
 	LOG_DIR := os.Getenv("LOG_DIR")
 	if LOG_DIR == "" {
-		LOG_DIR = "/var/coverage/data"
+		LOG_DIR = defaultLogDir
 	}
 	SAFE_BIN_DIR := os.Getenv("SAFE_BIN_DIR")
 	if SAFE_BIN_DIR == "" {
-		SAFE_BIN_DIR = "/var/coverage/bin"
+		SAFE_BIN_DIR = defaultSafeBinDir
 	}
 	pinTool, err := findPinTool(PIN_TOOL_SEARCH_DIR)
 	if err != nil {
@@ -110,7 +115,7 @@ func wrap(targetBinary string) error {
 	if err != nil {
 		return err
 	}
-	// change tmpdir permissions to allow execution from all users 
+	// change tmpdir permissions to allow execution from all users
 	if err := os.Chmod(tmpDir, 0755); err != nil {
 		return fmt.Errorf("could not set permissions on temp dir: %w", err)
 	}
@@ -184,7 +189,35 @@ func findPinTool(searchDir string) (string, error) {
 		return nil
 	})
 	if found == "" {
-		return "", errors.New("FuncTracer.so not found")
+		return "", errors.New("FuncTracer.so not found. Look for it in the $PIN_TOOL_SEARCH_DIR env variable or " + defaultPinToolSearchDir + " directory")
 	}
 	return found, nil
 }
+
+func wrapMany(binaries []string) error {
+	var failed []string
+	for _, bin := range binaries {
+		if err := wrap(bin); err != nil {
+			fmt.Fprintf(os.Stderr, "wrap error for %s: %v\n", bin, err)
+			failed = append(failed, bin)
+		}
+	}
+	if len(failed) > 0 {
+		return fmt.Errorf("failed to wrap: %v", failed)
+	}
+	return nil
+}
+
+func unwrapMany(binaries []string) error {
+	var failed []string
+	for _, bin := range binaries {
+		if err := unwrap(bin); err != nil {
+			fmt.Fprintf(os.Stderr, "unwrap error for %s: %v\n", bin, err)
+			failed = append(failed, bin)
+		}
+	}
+	if len(failed) > 0 {
+		return fmt.Errorf("failed to unwrap: %v", failed)
+	}
+	return nil
+}