-
Notifications
You must be signed in to change notification settings - Fork 2k
Expand file tree
/
Copy pathtrapwriter.go
More file actions
176 lines (162 loc) · 4.29 KB
/
trapwriter.go
File metadata and controls
176 lines (162 loc) · 4.29 KB
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
package trap
import (
"bufio"
"compress/gzip"
"errors"
"fmt"
"go/ast"
"go/types"
"os"
"path/filepath"
"unicode/utf8"
"github.com/github/codeql-go/extractor/srcarchive"
"golang.org/x/tools/go/packages"
)
// A Writer provides methods for writing data to a TRAP file
type Writer struct {
zip *gzip.Writer
wzip *bufio.Writer
wfile *bufio.Writer
file *os.File
Labeler *Labeler
path string
trapFilePath string
Package *packages.Package
TypesOverride map[ast.Expr]types.Type
ObjectsOverride map[types.Object]types.Object
}
func FileFor(path string) (string, error) {
trapFolder, err := trapFolder()
if err != nil {
return "", err
}
return filepath.Join(trapFolder, srcarchive.AppendablePath(path)+".trap.gz"), nil
}
// NewWriter creates a TRAP file for the given path and returns a writer for
// writing to it
func NewWriter(path string, pkg *packages.Package) (*Writer, error) {
trapFilePath, err := FileFor(path)
if err != nil {
return nil, err
}
trapFileDir := filepath.Dir(trapFilePath)
err = os.MkdirAll(trapFileDir, 0755)
if err != nil {
return nil, err
}
tmpFile, err := os.CreateTemp(trapFileDir, filepath.Base(trapFilePath))
if err != nil {
return nil, err
}
bufioFileWriter := bufio.NewWriter(tmpFile)
zipWriter := gzip.NewWriter(bufioFileWriter)
bufioZipWriter := bufio.NewWriter(zipWriter)
tw := &Writer{
zipWriter,
bufioZipWriter,
bufioFileWriter,
tmpFile,
nil,
path,
trapFilePath,
pkg,
make(map[ast.Expr]types.Type),
make(map[types.Object]types.Object),
}
tw.Labeler = newLabeler(tw)
return tw, nil
}
func trapFolder() (string, error) {
trapFolder := os.Getenv("CODEQL_EXTRACTOR_GO_TRAP_DIR")
if trapFolder == "" {
trapFolder = os.Getenv("TRAP_FOLDER")
}
if trapFolder == "" {
return "", errors.New("environment variable CODEQL_EXTRACTOR_GO_TRAP_DIR not set")
}
err := os.MkdirAll(trapFolder, 0755)
if err != nil {
return "", err
}
return trapFolder, nil
}
// Close the underlying file writer
func (tw *Writer) Close() error {
err := tw.wzip.Flush()
if err != nil {
// throw away file close error
tw.file.Close()
return err
}
err = tw.zip.Close()
if err != nil {
// return zip-close error, but ignore file-close error
tw.file.Close()
return err
}
err = tw.wfile.Flush()
if err != nil {
// throw away close error because write errors are likely to be more important
tw.file.Close()
return err
}
err = tw.file.Close()
if err != nil {
return err
}
return os.Rename(tw.file.Name(), tw.trapFilePath)
}
// ForEachObject iterates over all objects labeled by this labeler, and invokes
// the provided callback with a writer for the trap file, the object, and its
// label. It returns true if any extra objects were labeled and false otherwise.
func (tw *Writer) ForEachObject(cb func(*Writer, types.Object, Label)) bool {
// copy the objects into an array so that our behaviour is deterministic even
// if `cb` adds any new objects
i := 0
objects := make([]types.Object, len(tw.Labeler.objectLabels))
for k := range tw.Labeler.objectLabels {
objects[i] = k
i++
}
for _, object := range objects {
cb(tw, object, tw.Labeler.objectLabels[object])
}
return len(tw.Labeler.objectLabels) != len(objects)
}
const max_strlen = 1024 * 1024
func capStringLength(s string) string {
// if the UTF8-encoded string is longer than 1MiB, we truncate it
if len(s) > max_strlen {
// to ensure that the truncated string is valid UTF-8, we find the last byte at or
// before index max_strlen that starts a UTF-8 encoded character, and then cut off
// right before that byte
end := max_strlen
for ; !utf8.RuneStart(s[end]); end-- {
}
return s[0:end]
}
return s
}
// Emit writes out a tuple of values for the given `table`
func (tw *Writer) Emit(table string, values []interface{}) error {
fmt.Fprintf(tw.wzip, "%s(", table)
for i, value := range values {
if i > 0 {
fmt.Fprint(tw.wzip, ", ")
}
switch value := value.(type) {
case Label:
fmt.Fprint(tw.wzip, value.id)
case string:
fmt.Fprintf(tw.wzip, "\"%s\"", escapeString(capStringLength(value)))
case int:
fmt.Fprintf(tw.wzip, "%d", value)
case float64:
fmt.Fprintf(tw.wzip, "%e", value)
default:
return errors.New("Cannot emit value")
}
}
fmt.Fprintf(tw.wzip, ")\n")
return nil
}