forked from SoveraNia/SyzVegas
/
workqueue.go
218 lines (199 loc) · 5.18 KB
/
workqueue.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
// Copyright 2017 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.
// MODIFIED: Daimeng Wang
package main
import (
"sync"
"github.com/google/syzkaller/pkg/hash"
"github.com/google/syzkaller/pkg/ipc"
"github.com/google/syzkaller/prog"
)
// WorkQueue holds global non-fuzzing work items (see the Work* structs below).
// WorkQueue also does prioritization among work items, for example, we want
// to triage and send to manager new inputs before we smash programs
// in order to not permanently lose interesting programs in case of VM crash.
type WorkQueue struct {
mu sync.RWMutex
triageCandidate []*WorkTriage
candidate []*WorkCandidate
triage []*WorkTriage
smash []*WorkSmash
fuzzer *Fuzzer
procs int
needCandidates chan struct{}
}
type ProgTypes int
const (
ProgCandidate ProgTypes = 1 << iota
ProgMinimized
ProgSmashed
ProgNormal ProgTypes = 0
)
// WorkTriage are programs for which we noticed potential new coverage during
// first execution. But we are not sure yet if the coverage is real or not.
// During triage we understand if these programs in fact give new coverage,
// and if yes, minimize them and add to corpus.
type WorkTriage struct {
p *prog.Prog
call int
info ipc.CallInfo
flags ProgTypes
}
// WorkCandidate are programs from hub.
// We don't know yet if they are useful for this fuzzer or not.
// A proc handles them the same way as locally generated/mutated programs.
type WorkCandidate struct {
p *prog.Prog
flags ProgTypes
}
// WorkSmash are programs just added to corpus.
// During smashing these programs receive a one-time special attention
// (emit faults, collect comparison hints, etc).
type WorkSmash struct {
p *prog.Prog
call int
count int
}
func newWorkQueue(procs int, needCandidates chan struct{}) *WorkQueue {
return &WorkQueue{
procs: procs,
needCandidates: needCandidates,
}
}
func (wq *WorkQueue) enqueue(item interface{}) {
wq.mu.Lock()
defer wq.mu.Unlock()
switch item := item.(type) {
case *WorkTriage:
if item.flags&ProgCandidate != 0 {
wq.triageCandidate = append(wq.triageCandidate, item)
} else {
wq.triage = append(wq.triage, item)
}
case *WorkCandidate:
wq.candidate = append(wq.candidate, item)
case *WorkSmash:
wq.smash = append(wq.smash, item)
default:
panic("unknown work type")
}
}
func (wq *WorkQueue) dequeueType(t int, reverse bool, lock bool) (item interface{}) {
// 0 = candidate, 1 = smash, 2 = triage
// reverse = true -> LIFO
// Note: this should be done when lock is obtained
// ts0 := time.Now().UnixNano()
if lock {
wq.mu.Lock()
defer wq.mu.Unlock()
}
item = nil
switch t {
case 0:
{
if len(wq.candidate) != 0 {
if reverse {
last := len(wq.candidate) - 1
item = wq.candidate[last]
wq.candidate = wq.candidate[:last]
} else {
item = wq.candidate[0]
wq.candidate = wq.candidate[1:]
}
}
}
case 1:
{
if len(wq.smash) != 0 {
if reverse {
last := len(wq.smash) - 1
item = wq.smash[last]
wq.smash = wq.smash[:last]
} else {
item = wq.smash[0]
wq.smash = wq.smash[1:]
}
}
}
case 2:
{
if len(wq.triageCandidate) != 0 {
if reverse {
last := len(wq.triageCandidate) - 1
item = wq.triageCandidate[last]
wq.triageCandidate = wq.triageCandidate[:last]
} else {
item = wq.triageCandidate[0]
wq.triageCandidate = wq.triageCandidate[1:]
}
} else if len(wq.triage) != 0 {
if reverse {
last := len(wq.triage) - 1
item = wq.triage[last]
wq.triage = wq.triage[:last]
} else {
item = wq.triage[0]
wq.triage = wq.triage[1:]
}
}
}
}
// ts1 := time.Now().UnixNano()
// wq.fuzzer.writeLog("- MAB Dequeue: %v\n", ts1-ts0)
return item
}
func (wq *WorkQueue) dequeue() (item interface{}) {
wq.mu.RLock()
if len(wq.triageCandidate)+len(wq.candidate)+len(wq.triage)+len(wq.smash) == 0 {
wq.mu.RUnlock()
return nil
}
wq.mu.RUnlock()
wq.mu.Lock()
wantCandidates := false
if len(wq.triageCandidate) != 0 {
last := len(wq.triageCandidate) - 1
item = wq.triageCandidate[last]
wq.triageCandidate = wq.triageCandidate[:last]
} else if len(wq.candidate) != 0 {
last := len(wq.candidate) - 1
item = wq.candidate[last]
wq.candidate = wq.candidate[:last]
wantCandidates = len(wq.candidate) < wq.procs
} else if len(wq.triage) != 0 {
last := len(wq.triage) - 1
item = wq.triage[last]
wq.triage = wq.triage[:last]
} else if len(wq.smash) != 0 {
last := len(wq.smash) - 1
item = wq.smash[last]
wq.smash = wq.smash[:last]
}
wq.mu.Unlock()
if wantCandidates {
select {
case wq.needCandidates <- struct{}{}:
default:
}
}
return item
}
func (wq *WorkQueue) wantCandidates() bool {
wq.mu.RLock()
defer wq.mu.RUnlock()
return len(wq.candidate) < wq.procs
}
func (wq *WorkQueue) wantTriages() bool {
wq.mu.RLock()
defer wq.mu.RUnlock()
return (len(wq.triage) + len(wq.triageCandidate)) < wq.procs
}
func (wq *WorkQueue) findTriageProg(sig hash.Sig) *prog.Prog {
for _, item := range wq.triage {
_sig := hash.Hash(item.p.Serialize())
if sig == _sig {
return item.p
}
}
return nil
}