1
package tracker
2

3
import (
4
	"bytes"
5
	"fmt"
6
	"sync"
7

8
	"github.com/pelletier/go-toml/v2/unstable"
9
)
10

11
type keyKind uint8
12

13
const (
14
	invalidKind keyKind = iota
15
	valueKind
16
	tableKind
17
	arrayTableKind
18
)
19

20
func (k keyKind) String() string {
21
	switch k {
22
	case invalidKind:
23
		return "invalid"
24
	case valueKind:
25
		return "value"
26
	case tableKind:
27
		return "table"
28
	case arrayTableKind:
29
		return "array table"
30
	}
31
	panic("missing keyKind string mapping")
32
}
33

34
// SeenTracker tracks which keys have been seen with which TOML type to flag
35
// duplicates and mismatches according to the spec.
36
//
37
// Each node in the visited tree is represented by an entry. Each entry has an
38
// identifier, which is provided by a counter. Entries are stored in the array
39
// entries. As new nodes are discovered (referenced for the first time in the
40
// TOML document), entries are created and appended to the array. An entry
41
// points to its parent using its id.
42
//
43
// To find whether a given key (sequence of []byte) has already been visited,
44
// the entries are linearly searched, looking for one with the right name and
45
// parent id.
46
//
47
// Given that all keys appear in the document after their parent, it is
48
// guaranteed that all descendants of a node are stored after the node, this
49
// speeds up the search process.
50
//
51
// When encountering [[array tables]], the descendants of that node are removed
52
// to allow that branch of the tree to be "rediscovered". To maintain the
53
// invariant above, the deletion process needs to keep the order of entries.
54
// This results in more copies in that case.
55
type SeenTracker struct {
56
	entries    []entry
57
	currentIdx int
58
}
59

60
var pool = sync.Pool{
61
	New: func() interface{} {
62
		return &SeenTracker{}
63
	},
64
}
65

66
func (s *SeenTracker) reset() {
67
	// Always contains a root element at index 0.
68
	s.currentIdx = 0
69
	if len(s.entries) == 0 {
70
		s.entries = make([]entry, 1, 2)
71
	} else {
72
		s.entries = s.entries[:1]
73
	}
74
	s.entries[0].child = -1
75
	s.entries[0].next = -1
76
}
77

78
type entry struct {
79
	// Use -1 to indicate no child or no sibling.
80
	child int
81
	next  int
82

83
	name     []byte
84
	kind     keyKind
85
	explicit bool
86
	kv       bool
87
}
88

89
// Find the index of the child of parentIdx with key k. Returns -1 if
90
// it does not exist.
91
func (s *SeenTracker) find(parentIdx int, k []byte) int {
92
	for i := s.entries[parentIdx].child; i >= 0; i = s.entries[i].next {
93
		if bytes.Equal(s.entries[i].name, k) {
94
			return i
95
		}
96
	}
97
	return -1
98
}
99

100
// Remove all descendants of node at position idx.
101
func (s *SeenTracker) clear(idx int) {
102
	if idx >= len(s.entries) {
103
		return
104
	}
105

106
	for i := s.entries[idx].child; i >= 0; {
107
		next := s.entries[i].next
108
		n := s.entries[0].next
109
		s.entries[0].next = i
110
		s.entries[i].next = n
111
		s.entries[i].name = nil
112
		s.clear(i)
113
		i = next
114
	}
115

116
	s.entries[idx].child = -1
117
}
118

119
func (s *SeenTracker) create(parentIdx int, name []byte, kind keyKind, explicit bool, kv bool) int {
120
	e := entry{
121
		child: -1,
122
		next:  s.entries[parentIdx].child,
123

124
		name:     name,
125
		kind:     kind,
126
		explicit: explicit,
127
		kv:       kv,
128
	}
129
	var idx int
130
	if s.entries[0].next >= 0 {
131
		idx = s.entries[0].next
132
		s.entries[0].next = s.entries[idx].next
133
		s.entries[idx] = e
134
	} else {
135
		idx = len(s.entries)
136
		s.entries = append(s.entries, e)
137
	}
138

139
	s.entries[parentIdx].child = idx
140

141
	return idx
142
}
143

144
func (s *SeenTracker) setExplicitFlag(parentIdx int) {
145
	for i := s.entries[parentIdx].child; i >= 0; i = s.entries[i].next {
146
		if s.entries[i].kv {
147
			s.entries[i].explicit = true
148
			s.entries[i].kv = false
149
		}
150
		s.setExplicitFlag(i)
151
	}
152
}
153

154
// CheckExpression takes a top-level node and checks that it does not contain
155
// keys that have been seen in previous calls, and validates that types are
156
// consistent. It returns true if it is the first time this node's key is seen.
157
// Useful to clear array tables on first use.
158
func (s *SeenTracker) CheckExpression(node *unstable.Node) (bool, error) {
159
	if s.entries == nil {
160
		s.reset()
161
	}
162
	switch node.Kind {
163
	case unstable.KeyValue:
164
		return s.checkKeyValue(node)
165
	case unstable.Table:
166
		return s.checkTable(node)
167
	case unstable.ArrayTable:
168
		return s.checkArrayTable(node)
169
	default:
170
		panic(fmt.Errorf("this should not be a top level node type: %s", node.Kind))
171
	}
172
}
173

174
func (s *SeenTracker) checkTable(node *unstable.Node) (bool, error) {
175
	if s.currentIdx >= 0 {
176
		s.setExplicitFlag(s.currentIdx)
177
	}
178

179
	it := node.Key()
180

181
	parentIdx := 0
182

183
	// This code is duplicated in checkArrayTable. This is because factoring
184
	// it in a function requires to copy the iterator, or allocate it to the
185
	// heap, which is not cheap.
186
	for it.Next() {
187
		if it.IsLast() {
188
			break
189
		}
190

191
		k := it.Node().Data
192

193
		idx := s.find(parentIdx, k)
194

195
		if idx < 0 {
196
			idx = s.create(parentIdx, k, tableKind, false, false)
197
		} else {
198
			entry := s.entries[idx]
199
			if entry.kind == valueKind {
200
				return false, fmt.Errorf("toml: expected %s to be a table, not a %s", string(k), entry.kind)
201
			}
202
		}
203
		parentIdx = idx
204
	}
205

206
	k := it.Node().Data
207
	idx := s.find(parentIdx, k)
208

209
	first := false
210
	if idx >= 0 {
211
		kind := s.entries[idx].kind
212
		if kind != tableKind {
213
			return false, fmt.Errorf("toml: key %s should be a table, not a %s", string(k), kind)
214
		}
215
		if s.entries[idx].explicit {
216
			return false, fmt.Errorf("toml: table %s already exists", string(k))
217
		}
218
		s.entries[idx].explicit = true
219
	} else {
220
		idx = s.create(parentIdx, k, tableKind, true, false)
221
		first = true
222
	}
223

224
	s.currentIdx = idx
225

226
	return first, nil
227
}
228

229
func (s *SeenTracker) checkArrayTable(node *unstable.Node) (bool, error) {
230
	if s.currentIdx >= 0 {
231
		s.setExplicitFlag(s.currentIdx)
232
	}
233

234
	it := node.Key()
235

236
	parentIdx := 0
237

238
	for it.Next() {
239
		if it.IsLast() {
240
			break
241
		}
242

243
		k := it.Node().Data
244

245
		idx := s.find(parentIdx, k)
246

247
		if idx < 0 {
248
			idx = s.create(parentIdx, k, tableKind, false, false)
249
		} else {
250
			entry := s.entries[idx]
251
			if entry.kind == valueKind {
252
				return false, fmt.Errorf("toml: expected %s to be a table, not a %s", string(k), entry.kind)
253
			}
254
		}
255

256
		parentIdx = idx
257
	}
258

259
	k := it.Node().Data
260
	idx := s.find(parentIdx, k)
261

262
	firstTime := idx < 0
263
	if firstTime {
264
		idx = s.create(parentIdx, k, arrayTableKind, true, false)
265
	} else {
266
		kind := s.entries[idx].kind
267
		if kind != arrayTableKind {
268
			return false, fmt.Errorf("toml: key %s already exists as a %s,  but should be an array table", kind, string(k))
269
		}
270
		s.clear(idx)
271
	}
272

273
	s.currentIdx = idx
274

275
	return firstTime, nil
276
}
277

278
func (s *SeenTracker) checkKeyValue(node *unstable.Node) (bool, error) {
279
	parentIdx := s.currentIdx
280
	it := node.Key()
281

282
	for it.Next() {
283
		k := it.Node().Data
284

285
		idx := s.find(parentIdx, k)
286

287
		if idx < 0 {
288
			idx = s.create(parentIdx, k, tableKind, false, true)
289
		} else {
290
			entry := s.entries[idx]
291
			if it.IsLast() {
292
				return false, fmt.Errorf("toml: key %s is already defined", string(k))
293
			} else if entry.kind != tableKind {
294
				return false, fmt.Errorf("toml: expected %s to be a table, not a %s", string(k), entry.kind)
295
			} else if entry.explicit {
296
				return false, fmt.Errorf("toml: cannot redefine table %s that has already been explicitly defined", string(k))
297
			}
298
		}
299

300
		parentIdx = idx
301
	}
302

303
	s.entries[parentIdx].kind = valueKind
304

305
	value := node.Value()
306

307
	switch value.Kind {
308
	case unstable.InlineTable:
309
		return s.checkInlineTable(value)
310
	case unstable.Array:
311
		return s.checkArray(value)
312
	}
313

314
	return false, nil
315
}
316

317
func (s *SeenTracker) checkArray(node *unstable.Node) (first bool, err error) {
318
	it := node.Children()
319
	for it.Next() {
320
		n := it.Node()
321
		switch n.Kind {
322
		case unstable.InlineTable:
323
			first, err = s.checkInlineTable(n)
324
			if err != nil {
325
				return false, err
326
			}
327
		case unstable.Array:
328
			first, err = s.checkArray(n)
329
			if err != nil {
330
				return false, err
331
			}
332
		}
333
	}
334
	return first, nil
335
}
336

337
func (s *SeenTracker) checkInlineTable(node *unstable.Node) (first bool, err error) {
338
	s = pool.Get().(*SeenTracker)
339
	s.reset()
340

341
	it := node.Children()
342
	for it.Next() {
343
		n := it.Node()
344
		first, err = s.checkKeyValue(n)
345
		if err != nil {
346
			return false, err
347
		}
348
	}
349

350
	// As inline tables are self-contained, the tracker does not
351
	// need to retain the details of what they contain. The
352
	// keyValue element that creates the inline table is kept to
353
	// mark the presence of the inline table and prevent
354
	// redefinition of its keys: check* functions cannot walk into
355
	// a value.
356
	pool.Put(s)
357
	return first, nil
358
}
359

360