Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: close CPU profile #416

Merged
merged 1 commit into from
May 22, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 10 additions & 4 deletions bin/bin.go
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ func dumpHeapProfile() {
// If CPU profiling is enabled (ZGRAB2_CPUPROFILE is not empty), start tracking
// CPU profiling in the configured file. Caller is responsible for invoking
// stopCPUProfile() when finished.
func startCPUProfile() {
func startCPUProfile() *os.File {
if file := getCPUProfileFile(); file != "" {
now := time.Now()
fullFile := getFormattedFile(file, now)
Expand All @@ -72,24 +72,30 @@ func startCPUProfile() {
if err := pprof.StartCPUProfile(f); err != nil {
log.Fatal("could not start CPU profile: ", err)
}
return f
}

return nil
}

// If CPU profiling is enabled (ZGRAB2_CPUPROFILE is not empty), stop profiling
// CPU usage.
func stopCPUProfile() {
func stopCPUProfile(f *os.File) {
if getCPUProfileFile() != "" {
pprof.StopCPUProfile()
}
if f != nil {
f.Close()
}
}

// ZGrab2Main should be called by func main() in a binary. The caller is
// responsible for importing any modules in use. This allows clients to easily
// include custom sets of scan modules by creating new main packages with custom
// sets of ZGrab modules imported with side-effects.
func ZGrab2Main() {
startCPUProfile()
defer stopCPUProfile()
f := startCPUProfile()
defer stopCPUProfile(f)
defer dumpHeapProfile()
_, moduleType, flag, err := zgrab2.ParseCommandLine(os.Args[1:])

Expand Down
Loading