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
|
// Copyright 2021 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 main
import (
"container/heap"
"sync"
"sync/atomic"
"time"
"github.com/google/syzkaller/pkg/rpctype"
"github.com/google/syzkaller/prog"
)
type EnvDescr int64
const (
AnyEnvironment EnvDescr = iota
NewEnvironment
// TODO: add CleanVMEnvironment support.
EnvironmentsCount
)
// ExecTask is the atomic analysis entity. Once executed, it could trigger the
// pipeline propagation fof the program.
type ExecTask struct {
CreationTime time.Time
Program *prog.Prog
ID int64
ExecResultChan ExecResultChan
priority int // The priority of the item in the queue.
// The index is needed by update and is maintained by the heap.Interface methods.
index int // The index of the item in the heap.
}
func (t *ExecTask) ToRPC() *rpctype.ExecTask {
return &rpctype.ExecTask{
Prog: t.Program.Serialize(),
ID: t.ID,
}
}
type ExecTaskFactory struct {
chanMapMutex sync.Mutex
taskIDToExecResultChan map[int64]ExecResultChan
taskCounter int64
}
func MakeExecTaskFactory() *ExecTaskFactory {
return &ExecTaskFactory{
taskIDToExecResultChan: make(map[int64]ExecResultChan),
taskCounter: -1,
}
}
type ExecResultChan chan *ExecResult
func (factory *ExecTaskFactory) MakeExecTask(prog *prog.Prog) *ExecTask {
task := &ExecTask{
CreationTime: time.Now(),
Program: prog,
ExecResultChan: make(ExecResultChan),
ID: atomic.AddInt64(&factory.taskCounter, 1),
}
factory.chanMapMutex.Lock()
defer factory.chanMapMutex.Unlock()
factory.taskIDToExecResultChan[task.ID] = task.ExecResultChan
return task
}
func (factory *ExecTaskFactory) ExecTasksQueued() int {
factory.chanMapMutex.Lock()
defer factory.chanMapMutex.Unlock()
return len(factory.taskIDToExecResultChan)
}
func (factory *ExecTaskFactory) DeleteExecTask(task *ExecTask) {
factory.chanMapMutex.Lock()
defer factory.chanMapMutex.Unlock()
delete(factory.taskIDToExecResultChan, task.ID)
}
func (factory *ExecTaskFactory) GetExecResultChan(taskID int64) ExecResultChan {
factory.chanMapMutex.Lock()
defer factory.chanMapMutex.Unlock()
return factory.taskIDToExecResultChan[taskID]
}
func MakeExecTaskQueue() *ExecTaskQueue {
return &ExecTaskQueue{
pq: make(ExecTaskPriorityQueue, 0),
}
}
// ExecTaskQueue respects the pq.priority. Internally it is a thread-safe PQ.
type ExecTaskQueue struct {
pq ExecTaskPriorityQueue
mu sync.Mutex
}
// PopTask return false if no tasks are available.
func (q *ExecTaskQueue) PopTask() (*ExecTask, bool) {
q.mu.Lock()
defer q.mu.Unlock()
if q.pq.Len() == 0 {
return nil, false
}
return heap.Pop(&q.pq).(*ExecTask), true
}
func (q *ExecTaskQueue) PushTask(task *ExecTask) {
q.mu.Lock()
defer q.mu.Unlock()
heap.Push(&q.pq, task)
}
func (q *ExecTaskQueue) Len() int {
q.mu.Lock()
defer q.mu.Unlock()
return q.pq.Len()
}
// ExecTaskPriorityQueue reused example from https://pkg.go.dev/container/heap
type ExecTaskPriorityQueue []*ExecTask
func (pq ExecTaskPriorityQueue) Len() int { return len(pq) }
func (pq ExecTaskPriorityQueue) Less(i, j int) bool {
// We want Pop to give us the highest, not lowest, priority so we use greater than here.
return pq[i].priority > pq[j].priority
}
func (pq ExecTaskPriorityQueue) Swap(i, j int) {
pq[i], pq[j] = pq[j], pq[i]
pq[i].index = i
pq[j].index = j
}
func (pq *ExecTaskPriorityQueue) Push(x interface{}) {
n := len(*pq)
item := x.(*ExecTask)
item.index = n
*pq = append(*pq, item)
}
func (pq *ExecTaskPriorityQueue) Pop() interface{} {
old := *pq
n := len(old)
item := old[n-1]
old[n-1] = nil // avoid memory leak
item.index = -1 // for safety
*pq = old[0 : n-1]
return item
}
|