From 8f276ef29583e363bb886170f2f424f2d2a0e244 Mon Sep 17 00:00:00 2001
From: Aleksandr Nogikh <nogikh@google.com>
Date: Thu, 30 Jan 2025 15:25:31 +0100
Subject: pkg/vminfo: gracefully handle context abortion

On context abortion, return a special error.
On the pkg/rpcserver side, recognize and process it.
---
 pkg/fuzzer/queue/queue.go  | 4 ++--
 pkg/rpcserver/rpcserver.go | 4 ++++
 pkg/vminfo/syscalls.go     | 5 ++++-
 pkg/vminfo/vminfo.go       | 8 +++++++-
 4 files changed, 17 insertions(+), 4 deletions(-)

(limited to 'pkg')

diff --git a/pkg/fuzzer/queue/queue.go b/pkg/fuzzer/queue/queue.go
index 6e17225e2..1b98d768e 100644
--- a/pkg/fuzzer/queue/queue.go
+++ b/pkg/fuzzer/queue/queue.go
@@ -100,14 +100,14 @@ func (r *Request) Done(res *Result) {
 	close(r.done)
 }
 
-var errContextAborted = errors.New("context closed while waiting the result")
+var ErrRequestAborted = errors.New("context closed while waiting the result")
 
 // Wait() blocks until we have the result.
 func (r *Request) Wait(ctx context.Context) *Result {
 	r.initChannel()
 	select {
 	case <-ctx.Done():
-		return &Result{Status: ExecFailure, Err: errContextAborted}
+		return &Result{Status: ExecFailure, Err: ErrRequestAborted}
 	case <-r.done:
 		return r.result
 	}
diff --git a/pkg/rpcserver/rpcserver.go b/pkg/rpcserver/rpcserver.go
index bcfcbdf41..d0e6a15f1 100644
--- a/pkg/rpcserver/rpcserver.go
+++ b/pkg/rpcserver/rpcserver.go
@@ -420,6 +420,10 @@ func (serv *server) runCheck(ctx context.Context, info *handshakeResult) error {
 		close(serv.cfg.machineCheckStarted)
 	}
 	enabledCalls, disabledCalls, features, checkErr := serv.checker.Run(ctx, info.Files, info.Features)
+	if checkErr == vminfo.ErrAborted {
+		return nil
+	}
+
 	enabledCalls, transitivelyDisabled := serv.target.TransitivelyEnabledCalls(enabledCalls)
 	// Note: need to print disbled syscalls before failing due to an error.
 	// This helps to debug "all system calls are disabled".
diff --git a/pkg/vminfo/syscalls.go b/pkg/vminfo/syscalls.go
index 9793cfee8..f4a6919a5 100644
--- a/pkg/vminfo/syscalls.go
+++ b/pkg/vminfo/syscalls.go
@@ -115,7 +115,10 @@ func (ctx *checkContext) do(fileInfos []*flatrpc.FileInfo, featureInfos []*flatr
 	globs := make(map[string][]string)
 	for _, req := range globReqs {
 		res := req.Wait(ctx.ctx)
-		if res.Status != queue.Success {
+		if res.Err == queue.ErrRequestAborted {
+			// Don't return an error on context cancellation.
+			return nil, nil, nil, nil
+		} else if res.Status != queue.Success {
 			return nil, nil, nil, fmt.Errorf("failed to execute glob: %w (%v)\n%s\n%s",
 				res.Err, res.Status, req.GlobPattern, res.Output)
 		}
diff --git a/pkg/vminfo/vminfo.go b/pkg/vminfo/vminfo.go
index acacdc9d5..41fadf59a 100644
--- a/pkg/vminfo/vminfo.go
+++ b/pkg/vminfo/vminfo.go
@@ -101,10 +101,16 @@ func (checker *Checker) MachineInfo(fileInfos []*flatrpc.FileInfo) ([]*KernelMod
 	return modules, info.Bytes(), nil
 }
 
+var ErrAborted = errors.New("aborted through the context")
+
 func (checker *Checker) Run(ctx context.Context, files []*flatrpc.FileInfo, featureInfos []*flatrpc.FeatureInfo) (
 	map[*prog.Syscall]bool, map[*prog.Syscall]string, Features, error) {
 	cc := newCheckContext(ctx, checker.cfg, checker.checker, checker.executor)
-	return cc.do(files, featureInfos)
+	enabled, disabled, features, err := cc.do(files, featureInfos)
+	if ctx.Err() != nil {
+		return nil, nil, nil, ErrAborted
+	}
+	return enabled, disabled, features, err
 }
 
 // Implementation of the queue.Source interface.
-- 
cgit mrf-deployment