From 086ab166bbbf17d3f0a16c8b27f1995a70908b35 Mon Sep 17 00:00:00 2001
From: Aleksandr Nogikh <nogikh@google.com>
Date: Tue, 14 May 2024 16:52:34 +0200
Subject: pkg/fuzzer/queue: retry inputs from crashed VMs

Mark some requests as Important. The Retry() layer will give them one
more chance even if they were not executed due to a VM crash.

For now, the only important requests are related to triage, candidates
and pkg/vminfo tests.

Add tests for retry.go.
---
 pkg/fuzzer/queue/queue.go      | 10 ++++++
 pkg/fuzzer/queue/retry.go      |  6 ++++
 pkg/fuzzer/queue/retry_test.go | 72 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 88 insertions(+)
 create mode 100644 pkg/fuzzer/queue/retry_test.go

(limited to 'pkg/fuzzer/queue')

diff --git a/pkg/fuzzer/queue/queue.go b/pkg/fuzzer/queue/queue.go
index 14d4ace78..48b912846 100644
--- a/pkg/fuzzer/queue/queue.go
+++ b/pkg/fuzzer/queue/queue.go
@@ -40,11 +40,16 @@ type Request struct {
 	BinaryFile string // If set, it's executed instead of Prog.
 	Repeat     int    // Repeats in addition to the first run.
 
+	// Important requests will be retried even from crashed VMs.
+	Important bool
+
 	// The callback will be called on request completion in the LIFO order.
 	// If it returns false, all further processing will be stopped.
 	// It allows wrappers to intercept Done() requests.
 	callback DoneCallback
 
+	onceCrashed bool
+
 	mu     sync.Mutex
 	result *Result
 	done   chan struct{}
@@ -91,6 +96,11 @@ func (r *Request) Wait(ctx context.Context) *Result {
 	}
 }
 
+// Risky() returns true if there's a substantial risk of the input crashing the VM.
+func (r *Request) Risky() bool {
+	return r.onceCrashed
+}
+
 func (r *Request) hash() hash.Sig {
 	buf := new(bytes.Buffer)
 	if r.ExecOpts != nil {
diff --git a/pkg/fuzzer/queue/retry.go b/pkg/fuzzer/queue/retry.go
index 0b2e02ba5..c59a2c048 100644
--- a/pkg/fuzzer/queue/retry.go
+++ b/pkg/fuzzer/queue/retry.go
@@ -33,5 +33,11 @@ func (r *retryer) done(req *Request, res *Result) bool {
 		r.pq.Submit(req)
 		return false
 	}
+	// Retry important requests from crashed VMs once.
+	if res.Status == Crashed && req.Important && !req.onceCrashed {
+		req.onceCrashed = true
+		r.pq.Submit(req)
+		return false
+	}
 	return true
 }
diff --git a/pkg/fuzzer/queue/retry_test.go b/pkg/fuzzer/queue/retry_test.go
new file mode 100644
index 000000000..8529779a2
--- /dev/null
+++ b/pkg/fuzzer/queue/retry_test.go
@@ -0,0 +1,72 @@
+// Copyright 2024 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 queue
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRetryerOnRestart(t *testing.T) {
+	q := Plain()
+	retryerObj := Retry(q)
+
+	q.Submit(&Request{Important: true})
+	q.Submit(&Request{Important: false})
+
+	// The requests must be retried forever.
+	req1 := retryerObj.Next()
+	req2 := retryerObj.Next()
+	for i := 0; i < 10; i++ {
+		req1.Done(&Result{Status: Restarted})
+		req2.Done(&Result{Status: Restarted})
+		assert.Equal(t, req1, retryerObj.Next())
+		assert.Equal(t, req2, retryerObj.Next())
+	}
+
+	// Once successful, requests should no longer appear.
+	req1.Done(&Result{Status: Success})
+	req2.Done(&Result{Status: Success})
+
+	assert.Equal(t, Success, req1.Wait(context.Background()).Status)
+	assert.Equal(t, Success, req2.Wait(context.Background()).Status)
+
+	assert.Nil(t, retryerObj.Next())
+	assert.Nil(t, retryerObj.Next())
+}
+
+func TestRetryerOnCrash(t *testing.T) {
+	q := Plain()
+	retryerObj := Retry(q)
+
+	// Unimportant requests will not be retried.
+	req := &Request{Important: false}
+	q.Submit(req)
+	assert.Equal(t, req, retryerObj.Next())
+	req.Done(&Result{Status: Crashed})
+	assert.Nil(t, retryerObj.Next())
+	assert.Equal(t, Crashed, req.Wait(context.Background()).Status)
+
+	// Important requests will be retried once.
+	req = &Request{Important: true}
+	q.Submit(req)
+	assert.Equal(t, req, retryerObj.Next())
+	req.Done(&Result{Status: Crashed})
+	assert.Equal(t, req, retryerObj.Next())
+	req.Done(&Result{Status: Success})
+	assert.Nil(t, retryerObj.Next())
+	assert.Equal(t, Success, req.Wait(context.Background()).Status)
+
+	// .. but not more than once.
+	req = &Request{Important: true}
+	q.Submit(req)
+	assert.Equal(t, req, retryerObj.Next())
+	req.Done(&Result{Status: Crashed})
+	assert.Equal(t, req, retryerObj.Next())
+	req.Done(&Result{Status: Crashed})
+	assert.Nil(t, retryerObj.Next())
+	assert.Equal(t, Crashed, req.Wait(context.Background()).Status)
+}
-- 
cgit mrf-deployment