1// Copyright 2021 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package fuzz
6
7import (
8	"context"
9	"errors"
10	"flag"
11	"fmt"
12	"internal/race"
13	"io"
14	"os"
15	"os/signal"
16	"reflect"
17	"strconv"
18	"testing"
19	"time"
20)
21
22var benchmarkWorkerFlag = flag.Bool("benchmarkworker", false, "")
23
24func TestMain(m *testing.M) {
25	flag.Parse()
26	if *benchmarkWorkerFlag {
27		runBenchmarkWorker()
28		return
29	}
30	os.Exit(m.Run())
31}
32
33func BenchmarkWorkerFuzzOverhead(b *testing.B) {
34	if race.Enabled {
35		b.Skip("TODO(48504): fix and re-enable")
36	}
37	origEnv := os.Getenv("GODEBUG")
38	defer func() { os.Setenv("GODEBUG", origEnv) }()
39	os.Setenv("GODEBUG", fmt.Sprintf("%s,fuzzseed=123", origEnv))
40
41	ws := &workerServer{
42		fuzzFn:     func(_ CorpusEntry) (time.Duration, error) { return time.Second, nil },
43		workerComm: workerComm{memMu: make(chan *sharedMem, 1)},
44	}
45
46	mem, err := sharedMemTempFile(workerSharedMemSize)
47	if err != nil {
48		b.Fatalf("failed to create temporary shared memory file: %s", err)
49	}
50	defer func() {
51		if err := mem.Close(); err != nil {
52			b.Error(err)
53		}
54	}()
55
56	initialVal := []any{make([]byte, 32)}
57	encodedVals := marshalCorpusFile(initialVal...)
58	mem.setValue(encodedVals)
59
60	ws.memMu <- mem
61
62	b.ResetTimer()
63	for i := 0; i < b.N; i++ {
64		ws.m = newMutator()
65		mem.setValue(encodedVals)
66		mem.header().count = 0
67
68		ws.fuzz(context.Background(), fuzzArgs{Limit: 1})
69	}
70}
71
72// BenchmarkWorkerPing acts as the coordinator and measures the time it takes
73// a worker to respond to N pings. This is a rough measure of our RPC latency.
74func BenchmarkWorkerPing(b *testing.B) {
75	if race.Enabled {
76		b.Skip("TODO(48504): fix and re-enable")
77	}
78	b.SetParallelism(1)
79	w := newWorkerForTest(b)
80	for i := 0; i < b.N; i++ {
81		if err := w.client.ping(context.Background()); err != nil {
82			b.Fatal(err)
83		}
84	}
85}
86
87// BenchmarkWorkerFuzz acts as the coordinator and measures the time it takes
88// a worker to mutate a given input and call a trivial fuzz function N times.
89func BenchmarkWorkerFuzz(b *testing.B) {
90	if race.Enabled {
91		b.Skip("TODO(48504): fix and re-enable")
92	}
93	b.SetParallelism(1)
94	w := newWorkerForTest(b)
95	entry := CorpusEntry{Values: []any{[]byte(nil)}}
96	entry.Data = marshalCorpusFile(entry.Values...)
97	for i := int64(0); i < int64(b.N); {
98		args := fuzzArgs{
99			Limit:   int64(b.N) - i,
100			Timeout: workerFuzzDuration,
101		}
102		_, resp, _, err := w.client.fuzz(context.Background(), entry, args)
103		if err != nil {
104			b.Fatal(err)
105		}
106		if resp.Err != "" {
107			b.Fatal(resp.Err)
108		}
109		if resp.Count == 0 {
110			b.Fatal("worker did not make progress")
111		}
112		i += resp.Count
113	}
114}
115
116// newWorkerForTest creates and starts a worker process for testing or
117// benchmarking. The worker process calls RunFuzzWorker, which responds to
118// RPC messages until it's stopped. The process is stopped and cleaned up
119// automatically when the test is done.
120func newWorkerForTest(tb testing.TB) *worker {
121	tb.Helper()
122	c, err := newCoordinator(CoordinateFuzzingOpts{
123		Types: []reflect.Type{reflect.TypeOf([]byte(nil))},
124		Log:   io.Discard,
125	})
126	if err != nil {
127		tb.Fatal(err)
128	}
129	dir := ""             // same as self
130	binPath := os.Args[0] // same as self
131	args := append(os.Args[1:], "-benchmarkworker")
132	env := os.Environ() // same as self
133	w, err := newWorker(c, dir, binPath, args, env)
134	if err != nil {
135		tb.Fatal(err)
136	}
137	tb.Cleanup(func() {
138		if err := w.cleanup(); err != nil {
139			tb.Error(err)
140		}
141	})
142	if err := w.startAndPing(context.Background()); err != nil {
143		tb.Fatal(err)
144	}
145	tb.Cleanup(func() {
146		if err := w.stop(); err != nil {
147			tb.Error(err)
148		}
149	})
150	return w
151}
152
153func runBenchmarkWorker() {
154	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
155	defer cancel()
156	fn := func(CorpusEntry) error { return nil }
157	if err := RunFuzzWorker(ctx, fn); err != nil && err != ctx.Err() {
158		panic(err)
159	}
160}
161
162func BenchmarkWorkerMinimize(b *testing.B) {
163	if race.Enabled {
164		b.Skip("TODO(48504): fix and re-enable")
165	}
166
167	ws := &workerServer{
168		workerComm: workerComm{memMu: make(chan *sharedMem, 1)},
169	}
170
171	mem, err := sharedMemTempFile(workerSharedMemSize)
172	if err != nil {
173		b.Fatalf("failed to create temporary shared memory file: %s", err)
174	}
175	defer func() {
176		if err := mem.Close(); err != nil {
177			b.Error(err)
178		}
179	}()
180	ws.memMu <- mem
181
182	bytes := make([]byte, 1024)
183	ctx := context.Background()
184	for sz := 1; sz <= len(bytes); sz <<= 1 {
185		sz := sz
186		input := []any{bytes[:sz]}
187		encodedVals := marshalCorpusFile(input...)
188		mem = <-ws.memMu
189		mem.setValue(encodedVals)
190		ws.memMu <- mem
191		b.Run(strconv.Itoa(sz), func(b *testing.B) {
192			i := 0
193			ws.fuzzFn = func(_ CorpusEntry) (time.Duration, error) {
194				if i == 0 {
195					i++
196					return time.Second, errors.New("initial failure for deflake")
197				}
198				return time.Second, nil
199			}
200			for i := 0; i < b.N; i++ {
201				b.SetBytes(int64(sz))
202				ws.minimize(ctx, minimizeArgs{})
203			}
204		})
205	}
206}
207