-
Notifications
You must be signed in to change notification settings - Fork 3
/
io.go
168 lines (136 loc) · 3.68 KB
/
io.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
// Copyright 2021-2024 Zenauth Ltd.
// SPDX-License-Identifier: Apache-2.0
package internal
import (
"bufio"
"bytes"
"errors"
"fmt"
"io"
"io/fs"
"unicode"
policyv1 "github.com/cerbos/cerbos/api/genpb/cerbos/policy/v1"
schemav1 "github.com/cerbos/cerbos/api/genpb/cerbos/schema/v1"
"github.com/ghodss/yaml"
"google.golang.org/protobuf/encoding/protojson"
"google.golang.org/protobuf/proto"
)
const (
bufSize = 1024 * 4 // 4KiB
maxFileSize = 1024 * 1024 * 4 // 4MiB
newline = '\n'
)
var (
jsonStart = []byte("{")
yamlSep = []byte("---")
yamlComment = []byte("#")
ErrMultipleYAMLDocs = errors.New("more than one YAML document detected")
)
func ReadPolicyFromFile(fsys fs.FS, path string) (*policyv1.Policy, error) {
f, err := fsys.Open(path)
if err != nil {
return nil, fmt.Errorf("failed to open %s: %w", path, err)
}
defer f.Close()
return ReadPolicy(f)
}
// ReadPolicy reads a policy from the given reader.
func ReadPolicy(src io.Reader) (*policyv1.Policy, error) {
policy := &policyv1.Policy{}
if err := ReadJSONOrYAML(src, policy); err != nil {
return nil, err
}
return policy, nil
}
func ReadJSONOrYAML(src io.Reader, dest proto.Message) error {
d := mkDecoder(io.LimitReader(src, maxFileSize))
return d.decode(dest)
}
func mkDecoder(src io.Reader) decoder {
buf := bufio.NewReaderSize(src, bufSize)
prelude, _ := buf.Peek(bufSize)
trimmed := bytes.TrimLeftFunc(prelude, unicode.IsSpace)
if bytes.HasPrefix(trimmed, jsonStart) {
return newJSONDecoder(buf)
}
return newYAMLDecoder(buf)
}
type decoder interface {
decode(dest proto.Message) error
}
type decoderFunc func(dest proto.Message) error
func (df decoderFunc) decode(dest proto.Message) error {
return df(dest)
}
func newJSONDecoder(src *bufio.Reader) decoderFunc {
return func(dest proto.Message) error {
jsonBytes, err := io.ReadAll(src)
if err != nil {
return err
}
if err := protojson.Unmarshal(jsonBytes, dest); err != nil {
return fmt.Errorf("failed to unmarshal JSON: %w", err)
}
return nil
}
}
func newYAMLDecoder(src *bufio.Reader) decoderFunc {
return func(dest proto.Message) error {
buf := new(bytes.Buffer)
numDocs := 0
s := bufio.NewScanner(src)
seenContent := false
for s.Scan() {
line := s.Bytes()
trimmedLine := bytes.TrimSpace(line)
// ignore comments
if bytes.HasPrefix(trimmedLine, yamlComment) {
continue
}
// ignore empty lines at the beginning of the file
if !seenContent && len(trimmedLine) == 0 {
continue
}
seenContent = true
if bytes.HasPrefix(line, yamlSep) {
numDocs++
if numDocs > 1 || (numDocs == 1 && buf.Len() > 0) {
return ErrMultipleYAMLDocs
}
}
if _, err := buf.Write(line); err != nil {
return fmt.Errorf("failed to buffer YAML data: %w", err)
}
_ = buf.WriteByte(newline)
}
if err := s.Err(); err != nil {
return fmt.Errorf("failed to read from source: %w", err)
}
jsonBytes, err := yaml.YAMLToJSON(buf.Bytes())
if err != nil {
return fmt.Errorf("failed to convert YAML to JSON: %w", err)
}
if err := protojson.Unmarshal(jsonBytes, dest); err != nil {
return fmt.Errorf("failed to unmarshal JSON: %w", err)
}
return nil
}
}
func ReadSchemaFromFile(fsys fs.FS, path string) (*schemav1.Schema, error) {
f, err := fsys.Open(path)
if err != nil {
return nil, fmt.Errorf("failed to open %s: %w", path, err)
}
defer f.Close()
return ReadSchema(f, path)
}
func ReadSchema(src io.Reader, id string) (*schemav1.Schema, error) {
def, err := io.ReadAll(src)
if err != nil {
return nil, fmt.Errorf("failed to read all bytes from reader: %w", err)
}
return &schemav1.Schema{
Id: id,
Definition: def,
}, nil
}