vm: func Run accepts context

It allows to use context as a single termination signal source.

3 files changed, 12 insertions, 4 deletions

diff --git a/pkg/build/netbsd.go b/pkg/build/netbsd.go
index 31ae6de51..21564e1aa 100644
--- a/pkg/build/netbsd.go
+++ b/pkg/build/netbsd.go
@@ -4,6 +4,7 @@
 package build
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"os"
@@ -155,7 +156,9 @@ func (ctx netbsd) copyKernelToDisk(targetArch, vmType, outputDir, kernel string)
 	}
 	commands = append(commands, "mknod /dev/vhci c 355 0")
 	commands = append(commands, "sync") // Run sync so that the copied image is stored properly.
-	_, rep, err := inst.Run(time.Minute, reporter, strings.Join(commands, ";"))
+	ctxTimeout, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+	_, rep, err := inst.Run(ctxTimeout, reporter, strings.Join(commands, ";"))
 	if err != nil {
 		return fmt.Errorf("error syncing the instance %w", err)
 	}
diff --git a/pkg/instance/execprog.go b/pkg/instance/execprog.go
index 4fd656fc5..83461d07e 100644
--- a/pkg/instance/execprog.go
+++ b/pkg/instance/execprog.go
@@ -4,6 +4,7 @@
 package instance
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"time"
@@ -123,7 +124,9 @@ func (inst *ExecProgInstance) runCommand(command string, duration time.Duration,
 	if inst.BeforeContextLen != 0 {
 		opts = append(opts, vm.OutputSize(inst.BeforeContextLen))
 	}
-	output, rep, err := inst.VMInstance.Run(duration, inst.reporter, command, opts...)
+	ctxTimeout, cancel := context.WithTimeout(context.Background(), duration)
+	defer cancel()
+	output, rep, err := inst.VMInstance.Run(ctxTimeout, inst.reporter, command, opts...)
 	if err != nil {
 		return nil, fmt.Errorf("failed to run command in VM: %w", err)
 	}
diff --git a/pkg/manager/diff.go b/pkg/manager/diff.go
index dd9c7eebc..61fc91731 100644
--- a/pkg/manager/diff.go
+++ b/pkg/manager/diff.go
@@ -535,8 +535,10 @@ func (kc *kernelContext) runInstance(ctx context.Context, inst *vm.Instance,
 		return nil, fmt.Errorf("failed to parse manager's address")
 	}
 	cmd := fmt.Sprintf("%v runner %v %v %v", executorBin, inst.Index(), host, port)
-	_, rep, err := inst.Run(kc.cfg.Timeouts.VMRunningTime, kc.reporter, cmd,
-		vm.ExitTimeout, vm.StopContext(ctx), vm.InjectExecuting(injectExec),
+	ctxTimeout, cancel := context.WithTimeout(ctx, kc.cfg.Timeouts.VMRunningTime)
+	defer cancel()
+	_, rep, err := inst.Run(ctxTimeout, kc.reporter, cmd, vm.ExitTimeout,
+		vm.InjectExecuting(injectExec),
 		vm.EarlyFinishCb(func() {
 			// Depending on the crash type and kernel config, fuzzing may continue
 			// running for several seconds even after kernel has printed a crash report.