Skip to content

Commit

Permalink
feat(capture): show output progress during capture (#41)
Browse files Browse the repository at this point in the history
* feat(capture): show otuput progress during capture

Signed-off-by: Alessio Greggi <ale_grey_91@hotmail.it>
Co-authored-by: ccoVeille <3875889+ccoVeille@users.noreply.github.com>
  • Loading branch information
alegrey91 and ccoVeille authored Aug 29, 2024
1 parent d44a475 commit bb99b44
Show file tree
Hide file tree
Showing 4 changed files with 51 additions and 12 deletions.
6 changes: 4 additions & 2 deletions cmd/capture.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import (

var functionSymbols string
var commandOutput bool
var commandError bool
var libbpfOutput bool
var save bool
var directory string
Expand All @@ -50,7 +51,7 @@ by passing the function name symbol and the binary args.
for _, functionSymbol := range functionSymbolList {
syscalls, err := captor.Capture(functionSymbol, args, captureOpts)
if err != nil {
fmt.Printf("error capturing syscall: %w", err)
fmt.Printf("error capturing syscall: %v", err)
}

saveOpts := writer.WriteOptions{
Expand All @@ -72,7 +73,8 @@ func init() {
captureCmd.Flags().StringVarP(&functionSymbols, "functions", "f", "", "Name of the function symbols to be traced")
captureCmd.MarkFlagRequired("functions")

captureCmd.Flags().BoolVarP(&commandOutput, "include-cmd-output", "c", false, "Include the executed command output")
captureCmd.Flags().BoolVarP(&commandOutput, "include-cmd-stdout", "c", false, "Include the executed command output")
captureCmd.Flags().BoolVarP(&commandError, "include-cmd-stderr", "e", false, "Include the executed command error")

captureCmd.Flags().BoolVarP(&libbpfOutput, "include-libbpf-output", "l", false, "Include the libbpf output")

Expand Down
5 changes: 3 additions & 2 deletions cmd/hunt.go
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ var huntCmd = &cobra.Command{
for _, functionSymbol := range symbolsOrigins.Symbols {
syscalls, err := captor.Capture(functionSymbol, captureArgs, opts)
if err != nil {
fmt.Printf("error capturing syscall: %w", err)
fmt.Printf("error capturing syscall: %v", err)
}

saveOpts := writer.WriteOptions{
Expand All @@ -94,7 +94,8 @@ func init() {
huntCmd.Flags().StringVarP(&harpoonFile, "file", "F", "harpoon-report.yml", "File with the result of analysis")
huntCmd.MarkFlagRequired("file")

huntCmd.Flags().BoolVarP(&commandOutput, "include-cmd-output", "c", false, "Include the executed command output")
huntCmd.Flags().BoolVarP(&commandOutput, "include-cmd-stdout", "c", false, "Include the executed command output")
huntCmd.Flags().BoolVarP(&commandError, "include-cmd-stderr", "e", false, "Include the executed command error")

huntCmd.Flags().BoolVarP(&libbpfOutput, "include-libbpf-output", "l", false, "Include the libbpf output")

Expand Down
15 changes: 14 additions & 1 deletion internal/ebpf/probesfacade/captor/capture.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ type event struct {

type CaptureOptions struct {
CommandOutput bool
CommandError bool
LibbpfOutput bool
}

Expand Down Expand Up @@ -114,7 +115,9 @@ func Capture(functionSymbol string, cmdArgs []string, opts CaptureOptions) ([]ui
// run args that we want to trace
var wg sync.WaitGroup
wg.Add(1)
go executor.Run(cmdArgs, opts.CommandOutput, &wg)
outputCh := make(chan string)
errorCh := make(chan string)
go executor.Run(cmdArgs, opts.CommandOutput, opts.CommandError, &wg, outputCh, errorCh)

var syscalls []uint32
go func() {
Expand All @@ -131,6 +134,16 @@ func Capture(functionSymbol string, cmdArgs []string, opts CaptureOptions) ([]ui
case lost := <-lostChannel:
fmt.Fprintf(os.Stderr, "lost %d data\n", lost)
return
case line, ok := <-outputCh:
if !ok {
break
}
fmt.Println(line)
case err, ok := <-errorCh:
if !ok {
break
}
fmt.Println(err)
}
}
}()
Expand Down
37 changes: 30 additions & 7 deletions internal/executor/exec.go
Original file line number Diff line number Diff line change
@@ -1,27 +1,50 @@
package executor

import (
"bufio"
"fmt"
"os/exec"
"sync"
)

// Run execute the command and wait for its end.
// The cmdOutput argument is used to print the command output.
func Run(cmd []string, cmdOutput bool, wg *sync.WaitGroup) {
func Run(cmd []string, cmdOutput, cmdError bool, wg *sync.WaitGroup, outputCh, errorCh chan<- string) {
defer func() {
wg.Done()
}()

command := exec.Command(cmd[0], cmd[1:]...)
output, _ := command.CombinedOutput()
if cmdOutput {
fmt.Println("----- BEGIN OF COMMAND OUTPUT -----")
fmt.Printf("%s", output)
fmt.Println("------ END OF COMMAND OUTPUT ------")
}
stdout, _ := command.StdoutPipe()
stderr, _ := command.StderrPipe()

//command.Wait()
command.Start()

go func() {
if cmdOutput {
scanner := bufio.NewScanner(stdout)
for scanner.Scan() {
outputCh <- scanner.Text()
}
if err := scanner.Err(); err != nil {
outputCh <- fmt.Sprintf("error: %v", err)
}
}
if cmdError {
scanner := bufio.NewScanner(stderr)
for scanner.Scan() {
errorCh <- scanner.Text()
}
if err := scanner.Err(); err != nil {
errorCh <- fmt.Sprintf("error: %v", err)
}
}
}()

command.Wait()
close(outputCh)
close(errorCh)
}

func Build(packagePath, outputFile string) (string, error) {
Expand Down

0 comments on commit bb99b44

Please # to comment.