pkg/cmdprof: merge into pkg/tool

cmdprof functionality seems to fit well into pkg/tool.

8 files changed, 73 insertions, 82 deletions

diff --git a/pkg/cmdprof/cmdprof.go b/pkg/cmdprof/cmdprof.go
deleted file mode 100644
index ab0396d16..000000000
--- a/pkg/cmdprof/cmdprof.go
+++ /dev/null
@@ -1,57 +0,0 @@
-// Copyright 2020 syzkaller project authors. All rights reserved.
-// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
-
-// Package cmdprof simplifies cpu/memory profiling for command line tools. Use as:
-//	flag.Parse()
-//	defer cmdprof.Install()()
-package cmdprof
-
-import (
-	"flag"
-	"fmt"
-	"os"
-	"runtime"
-	"runtime/pprof"
-)
-
-var (
-	flagCPUProfile = flag.String("cpuprofile", "", "write CPU profile to this file")
-	flagMEMProfile = flag.String("memprofile", "", "write memory profile to this file")
-)
-
-func Install() func() {
-	res := func() {}
-	failf := func(msg string, args ...interface{}) {
-		fmt.Fprintf(os.Stderr, msg+"\n", args...)
-		os.Exit(1)
-	}
-	if *flagCPUProfile != "" {
-		f, err := os.Create(*flagCPUProfile)
-		if err != nil {
-			failf("failed to create cpuprofile file: %v", err)
-		}
-		if err := pprof.StartCPUProfile(f); err != nil {
-			failf("failed to start cpu profile: %v", err)
-		}
-		res = func() {
-			pprof.StopCPUProfile()
-			f.Close()
-		}
-	}
-	if *flagMEMProfile != "" {
-		prev := res
-		res = func() {
-			prev()
-			f, err := os.Create(*flagMEMProfile)
-			if err != nil {
-				failf("failed to create memprofile file: %v", err)
-			}
-			defer f.Close()
-			runtime.GC()
-			if err := pprof.WriteHeapProfile(f); err != nil {
-				failf("failed to write mem profile: %v", err)
-			}
-		}
-	}
-	return res
-}
diff --git a/pkg/tool/cmdprof.go b/pkg/tool/cmdprof.go
new file mode 100644
index 000000000..196a0cd24
--- /dev/null
+++ b/pkg/tool/cmdprof.go
@@ -0,0 +1,44 @@
+// Copyright 2020 syzkaller project authors. All rights reserved.
+// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
+
+package tool
+
+import (
+	"os"
+	"runtime"
+	"runtime/pprof"
+)
+
+// installProfiling simplifies cpu/memory profiling for command line tools.
+func installProfiling(cpuprof, memprof string) func() {
+	res := func() {}
+	if cpuprof != "" {
+		f, err := os.Create(cpuprof)
+		if err != nil {
+			Failf("failed to create cpuprofile file: %v", err)
+		}
+		if err := pprof.StartCPUProfile(f); err != nil {
+			Failf("failed to start cpu profile: %v", err)
+		}
+		res = func() {
+			pprof.StopCPUProfile()
+			f.Close()
+		}
+	}
+	if memprof != "" {
+		prev := res
+		res = func() {
+			prev()
+			f, err := os.Create(memprof)
+			if err != nil {
+				Failf("failed to create memprofile file: %v", err)
+			}
+			defer f.Close()
+			runtime.GC()
+			if err := pprof.WriteHeapProfile(f); err != nil {
+				Failf("failed to write mem profile: %v", err)
+			}
+		}
+	}
+	return res
+}
diff --git a/pkg/tool/flags.go b/pkg/tool/flags.go
index 197494823..0763cb502 100644
--- a/pkg/tool/flags.go
+++ b/pkg/tool/flags.go
@@ -8,39 +8,31 @@ import (
 	"encoding/hex"
 	"flag"
 	"fmt"
-	"os"
 	"strings"
 
 	"github.com/google/syzkaller/pkg/log"
 )
 
-// ParseFlags parses command line flags with flag.Parse and then applies optional flags created with OptionalFlags.
-// This is intended for programmatic use only when we invoke older versions of binaries with new unsupported flags.
-func ParseFlags() {
-	if err := parseFlags(flag.CommandLine, os.Args[1:]); err != nil {
-		Fail(err)
-	}
-}
-
 type Flag struct {
 	Name  string
 	Value string
 }
 
 // OptionalFlags produces command line flag value that encapsulates the given flags as optional.
-// Use ParseFlags to support optional flags in the binary.
+// This is intended for programmatic use only when we invoke older versions of binaries with new unsupported flags.
+// Use tool.Init to support optional flags in the binary.
 // The format keeps flags reasonably readable ("-optional=foo=bar:baz=123"), not subject to accidental splitting
 // into multiple arguments due to spaces and supports bool/non-bool flags.
-func OptionalFlags(flags ...Flag) string {
+func OptionalFlags(flags []Flag) string {
 	return fmt.Sprintf("-%v=%v", optionalFlag, serializeFlags(flags))
 }
 
-func parseFlags(set *flag.FlagSet, args []string) error {
-	optional := set.String(optionalFlag, "", "optional flags for programmatic use only")
+func ParseFlags(set *flag.FlagSet, args []string) error {
+	flagOptional := set.String(optionalFlag, "", "optional flags for programmatic use only")
 	if err := set.Parse(args); err != nil {
 		return err
 	}
-	flags, err := deserializeFlags(*optional)
+	flags, err := deserializeFlags(*flagOptional)
 	if err != nil {
 		return err
 	}
diff --git a/pkg/tool/flags_test.go b/pkg/tool/flags_test.go
index 6d52a4780..d818c6582 100644
--- a/pkg/tool/flags_test.go
+++ b/pkg/tool/flags_test.go
@@ -27,7 +27,7 @@ func TestParseFlags(t *testing.T) {
 		{"", &Values{false, 1, "baz"}},
 		{"-foo -bar=2", &Values{true, 2, "baz"}},
 		{"-foo -bar=2 -qux", nil},
-		{"-foo -bar=2 " + OptionalFlags(Flag{"qux", ""}), &Values{true, 2, "baz"}},
+		{"-foo -bar=2 " + OptionalFlags([]Flag{{"qux", ""}}), &Values{true, 2, "baz"}},
 	}
 	for i, test := range tests {
 		t.Run(fmt.Sprint(i), func(t *testing.T) {
@@ -41,7 +41,7 @@ func TestParseFlags(t *testing.T) {
 			if args[0] == "" {
 				args = args[1:]
 			}
-			err := parseFlags(flags, args)
+			err := ParseFlags(flags, args)
 			if test.vals == nil {
 				if err == nil {
 					t.Fatalf("parsing did not fail")
diff --git a/pkg/tool/tool.go b/pkg/tool/tool.go
index bb1f436cc..690bfa4b8 100644
--- a/pkg/tool/tool.go
+++ b/pkg/tool/tool.go
@@ -5,10 +5,25 @@
 package tool
 
 import (
+	"flag"
 	"fmt"
 	"os"
 )
 
+// Init handles common tasks for command line tools:
+//  - invokes flag.Parse
+//  - adds support for optional flags (see OptionalFlags)
+//  - adds support for cpu/mem profiling (-cpuprofile/memprofile flags)
+// Use as defer tool.Init()().
+func Init() func() {
+	flagCPUProfile := flag.String("cpuprofile", "", "write CPU profile to this file")
+	flagMEMProfile := flag.String("memprofile", "", "write memory profile to this file")
+	if err := ParseFlags(flag.CommandLine, os.Args[1:]); err != nil {
+		Fail(err)
+	}
+	return installProfiling(*flagCPUProfile, *flagMEMProfile)
+}
+
 func Failf(msg string, args ...interface{}) {
 	fmt.Fprintf(os.Stderr, msg+"\n", args...)
 	os.Exit(1)
diff --git a/sys/syz-sysgen/sysgen.go b/sys/syz-sysgen/sysgen.go
index 62f36c3d3..b62e1dff2 100644
--- a/sys/syz-sysgen/sysgen.go
+++ b/sys/syz-sysgen/sysgen.go
@@ -18,11 +18,11 @@ import (
 	"text/template"
 
 	"github.com/google/syzkaller/pkg/ast"
-	"github.com/google/syzkaller/pkg/cmdprof"
 	"github.com/google/syzkaller/pkg/compiler"
 	"github.com/google/syzkaller/pkg/hash"
 	"github.com/google/syzkaller/pkg/osutil"
 	"github.com/google/syzkaller/pkg/serializer"
+	"github.com/google/syzkaller/pkg/tool"
 	"github.com/google/syzkaller/prog"
 	"github.com/google/syzkaller/sys/targets"
 )
@@ -60,8 +60,7 @@ var srcDir = flag.String("src", "", "path to root of syzkaller source dir")
 var outDir = flag.String("out", "", "path to out dir")
 
 func main() {
-	flag.Parse()
-	defer cmdprof.Install()()
+	defer tool.Init()()
 
 	var OSList []string
 	for OS := range targets.List {
diff --git a/tools/syz-check/check.go b/tools/syz-check/check.go
index de7145dfb..5961cb6e8 100644
--- a/tools/syz-check/check.go
+++ b/tools/syz-check/check.go
@@ -37,10 +37,10 @@ import (
 	"unsafe"
 
 	"github.com/google/syzkaller/pkg/ast"
-	"github.com/google/syzkaller/pkg/cmdprof"
 	"github.com/google/syzkaller/pkg/compiler"
 	"github.com/google/syzkaller/pkg/osutil"
 	"github.com/google/syzkaller/pkg/symbolizer"
+	"github.com/google/syzkaller/pkg/tool"
 	"github.com/google/syzkaller/prog"
 	"github.com/google/syzkaller/sys/targets"
 )
@@ -59,8 +59,7 @@ func main() {
 		fmt.Fprintf(os.Stderr, msg+"\n", args...)
 		os.Exit(1)
 	}
-	flag.Parse()
-	defer cmdprof.Install()()
+	defer tool.Init()()
 	var warnings []Warn
 	for arch, obj := range arches {
 		if *obj == "" {
diff --git a/tools/syz-cover/syz-cover.go b/tools/syz-cover/syz-cover.go
index d019bfb78..2808cdd96 100644
--- a/tools/syz-cover/syz-cover.go
+++ b/tools/syz-cover/syz-cover.go
@@ -27,9 +27,9 @@ import (
 	"strconv"
 	"strings"
 
-	"github.com/google/syzkaller/pkg/cmdprof"
 	"github.com/google/syzkaller/pkg/cover"
 	"github.com/google/syzkaller/pkg/osutil"
+	"github.com/google/syzkaller/pkg/tool"
 	"github.com/google/syzkaller/sys/targets"
 )
 
@@ -43,8 +43,7 @@ func main() {
 		flagKernelObj      = flag.String("kernel_obj", "", "path to kernel build/obj dir")
 		flagExport         = flag.String("csv", "", "export coverage data in csv format (optional)")
 	)
-	flag.Parse()
-	defer cmdprof.Install()()
+	defer tool.Init()()
 
 	if len(flag.Args()) == 0 {
 		fmt.Fprintf(os.Stderr, "usage: syz-cover [flags] rawcover.file\n")