forked from hashicorp/consul
/
prepared_query.go
398 lines (349 loc) · 12.2 KB
/
prepared_query.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
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
package state
import (
"fmt"
"regexp"
"github.com/hashicorp/consul/agent/consul/prepared_query"
"github.com/hashicorp/consul/agent/structs"
"github.com/hashicorp/go-memdb"
)
// preparedQueriesTableSchema returns a new table schema used for storing
// prepared queries.
func preparedQueriesTableSchema() *memdb.TableSchema {
return &memdb.TableSchema{
Name: "prepared-queries",
Indexes: map[string]*memdb.IndexSchema{
"id": &memdb.IndexSchema{
Name: "id",
AllowMissing: false,
Unique: true,
Indexer: &memdb.UUIDFieldIndex{
Field: "ID",
},
},
"name": &memdb.IndexSchema{
Name: "name",
AllowMissing: true,
Unique: true,
Indexer: &memdb.StringFieldIndex{
Field: "Name",
Lowercase: true,
},
},
"template": &memdb.IndexSchema{
Name: "template",
AllowMissing: true,
Unique: true,
Indexer: &PreparedQueryIndex{},
},
"session": &memdb.IndexSchema{
Name: "session",
AllowMissing: true,
Unique: false,
Indexer: &memdb.UUIDFieldIndex{
Field: "Session",
},
},
},
}
}
func init() {
registerSchema(preparedQueriesTableSchema)
}
// validUUID is used to check if a given string looks like a UUID
var validUUID = regexp.MustCompile(`(?i)^[\da-f]{8}-[\da-f]{4}-[\da-f]{4}-[\da-f]{4}-[\da-f]{12}$`)
// isUUID returns true if the given string is a valid UUID.
func isUUID(str string) bool {
const uuidLen = 36
if len(str) != uuidLen {
return false
}
return validUUID.MatchString(str)
}
// queryWrapper is an internal structure that is used to store a query alongside
// its compiled template, which can be nil.
type queryWrapper struct {
// We embed the PreparedQuery structure so that the UUID field indexer
// can see the ID directly.
*structs.PreparedQuery
// ct is the compiled template, or nil if the query isn't a template. The
// state store manages this and keeps it up to date every time the query
// changes.
ct *prepared_query.CompiledTemplate
}
// toPreparedQuery unwraps the internal form of a prepared query and returns
// the regular struct.
func toPreparedQuery(wrapped interface{}) *structs.PreparedQuery {
if wrapped == nil {
return nil
}
return wrapped.(*queryWrapper).PreparedQuery
}
// PreparedQueries is used to pull all the prepared queries from the snapshot.
func (s *Snapshot) PreparedQueries() (structs.PreparedQueries, error) {
queries, err := s.tx.Get("prepared-queries", "id")
if err != nil {
return nil, err
}
var ret structs.PreparedQueries
for wrapped := queries.Next(); wrapped != nil; wrapped = queries.Next() {
ret = append(ret, toPreparedQuery(wrapped))
}
return ret, nil
}
// PreparedQuery is used when restoring from a snapshot. For general inserts,
// use PreparedQuerySet.
func (s *Restore) PreparedQuery(query *structs.PreparedQuery) error {
// If this is a template, compile it, otherwise leave the compiled
// template field nil.
var ct *prepared_query.CompiledTemplate
if prepared_query.IsTemplate(query) {
var err error
ct, err = prepared_query.Compile(query)
if err != nil {
return fmt.Errorf("failed compiling template: %s", err)
}
}
// Insert the wrapped query.
if err := s.tx.Insert("prepared-queries", &queryWrapper{query, ct}); err != nil {
return fmt.Errorf("failed restoring prepared query: %s", err)
}
if err := indexUpdateMaxTxn(s.tx, query.ModifyIndex, "prepared-queries"); err != nil {
return fmt.Errorf("failed updating index: %s", err)
}
return nil
}
// PreparedQuerySet is used to create or update a prepared query.
func (s *Store) PreparedQuerySet(idx uint64, query *structs.PreparedQuery) error {
tx := s.db.Txn(true)
defer tx.Abort()
if err := s.preparedQuerySetTxn(tx, idx, query); err != nil {
return err
}
tx.Commit()
return nil
}
// preparedQuerySetTxn is the inner method used to insert a prepared query with
// the proper indexes into the state store.
func (s *Store) preparedQuerySetTxn(tx *memdb.Txn, idx uint64, query *structs.PreparedQuery) error {
// Check that the ID is set.
if query.ID == "" {
return ErrMissingQueryID
}
// Check for an existing query.
wrapped, err := tx.First("prepared-queries", "id", query.ID)
if err != nil {
return fmt.Errorf("failed prepared query lookup: %s", err)
}
existing := toPreparedQuery(wrapped)
// Set the indexes.
if existing != nil {
query.CreateIndex = existing.CreateIndex
query.ModifyIndex = idx
} else {
query.CreateIndex = idx
query.ModifyIndex = idx
}
// Verify that the query name doesn't already exist, or that we are
// updating the same instance that has this name. If this is a template
// and the name is empty then we make sure there's not an empty template
// already registered.
if query.Name != "" {
wrapped, err := tx.First("prepared-queries", "name", query.Name)
if err != nil {
return fmt.Errorf("failed prepared query lookup: %s", err)
}
other := toPreparedQuery(wrapped)
if other != nil && (existing == nil || existing.ID != other.ID) {
return fmt.Errorf("name '%s' aliases an existing query name", query.Name)
}
} else if prepared_query.IsTemplate(query) {
wrapped, err := tx.First("prepared-queries", "template", query.Name)
if err != nil {
return fmt.Errorf("failed prepared query lookup: %s", err)
}
other := toPreparedQuery(wrapped)
if other != nil && (existing == nil || existing.ID != other.ID) {
return fmt.Errorf("a query template with an empty name already exists")
}
}
// Verify that the name doesn't alias any existing ID. We allow queries
// to be looked up by ID *or* name so we don't want anyone to try to
// register a query with a name equal to some other query's ID in an
// attempt to hijack it. We also look up by ID *then* name in order to
// prevent this, but it seems prudent to prevent these types of rogue
// queries from ever making it into the state store. Note that we have
// to see if the name looks like a UUID before checking since the UUID
// index will complain if we look up something that's not formatted
// like one.
if isUUID(query.Name) {
wrapped, err := tx.First("prepared-queries", "id", query.Name)
if err != nil {
return fmt.Errorf("failed prepared query lookup: %s", err)
}
if wrapped != nil {
return fmt.Errorf("name '%s' aliases an existing query ID", query.Name)
}
}
// Verify that the session exists.
if query.Session != "" {
sess, err := tx.First("sessions", "id", query.Session)
if err != nil {
return fmt.Errorf("failed session lookup: %s", err)
}
if sess == nil {
return fmt.Errorf("invalid session %#v", query.Session)
}
}
// We do not verify the service here, nor the token, if any. These are
// checked at execute time and not doing integrity checking on them
// helps avoid bootstrapping chicken and egg problems.
// If this is a template, compile it, otherwise leave the compiled
// template field nil.
var ct *prepared_query.CompiledTemplate
if prepared_query.IsTemplate(query) {
var err error
ct, err = prepared_query.Compile(query)
if err != nil {
return fmt.Errorf("failed compiling template: %s", err)
}
}
// Insert the wrapped query.
if err := tx.Insert("prepared-queries", &queryWrapper{query, ct}); err != nil {
return fmt.Errorf("failed inserting prepared query: %s", err)
}
if err := tx.Insert("index", &IndexEntry{"prepared-queries", idx}); err != nil {
return fmt.Errorf("failed updating index: %s", err)
}
return nil
}
// PreparedQueryDelete deletes the given query by ID.
func (s *Store) PreparedQueryDelete(idx uint64, queryID string) error {
tx := s.db.Txn(true)
defer tx.Abort()
if err := s.preparedQueryDeleteTxn(tx, idx, queryID); err != nil {
return fmt.Errorf("failed prepared query delete: %s", err)
}
tx.Commit()
return nil
}
// preparedQueryDeleteTxn is the inner method used to delete a prepared query
// with the proper indexes into the state store.
func (s *Store) preparedQueryDeleteTxn(tx *memdb.Txn, idx uint64, queryID string) error {
// Pull the query.
wrapped, err := tx.First("prepared-queries", "id", queryID)
if err != nil {
return fmt.Errorf("failed prepared query lookup: %s", err)
}
if wrapped == nil {
return nil
}
// Delete the query and update the index.
if err := tx.Delete("prepared-queries", wrapped); err != nil {
return fmt.Errorf("failed prepared query delete: %s", err)
}
if err := tx.Insert("index", &IndexEntry{"prepared-queries", idx}); err != nil {
return fmt.Errorf("failed updating index: %s", err)
}
return nil
}
// PreparedQueryGet returns the given prepared query by ID.
func (s *Store) PreparedQueryGet(ws memdb.WatchSet, queryID string) (uint64, *structs.PreparedQuery, error) {
tx := s.db.Txn(false)
defer tx.Abort()
// Get the table index.
idx := maxIndexTxn(tx, "prepared-queries")
// Look up the query by its ID.
watchCh, wrapped, err := tx.FirstWatch("prepared-queries", "id", queryID)
if err != nil {
return 0, nil, fmt.Errorf("failed prepared query lookup: %s", err)
}
ws.Add(watchCh)
return idx, toPreparedQuery(wrapped), nil
}
// PreparedQueryResolve returns the given prepared query by looking up an ID or
// Name. If the query was looked up by name and it's a template, then the
// template will be rendered before it is returned.
func (s *Store) PreparedQueryResolve(queryIDOrName string, source structs.QuerySource) (uint64, *structs.PreparedQuery, error) {
tx := s.db.Txn(false)
defer tx.Abort()
// Get the table index.
idx := maxIndexTxn(tx, "prepared-queries")
// Explicitly ban an empty query. This will never match an ID and the
// schema is set up so it will never match a query with an empty name,
// but we check it here to be explicit about it (we'd never want to
// return the results from the first query w/o a name).
if queryIDOrName == "" {
return 0, nil, ErrMissingQueryID
}
// Try first by ID if it looks like they gave us an ID. We check the
// format before trying this because the UUID index will complain if
// we look up something that's not formatted like one.
if isUUID(queryIDOrName) {
wrapped, err := tx.First("prepared-queries", "id", queryIDOrName)
if err != nil {
return 0, nil, fmt.Errorf("failed prepared query lookup: %s", err)
}
if wrapped != nil {
query := toPreparedQuery(wrapped)
if prepared_query.IsTemplate(query) {
return idx, nil, fmt.Errorf("prepared query templates can only be resolved up by name, not by ID")
}
return idx, query, nil
}
}
// prep will check to see if the query is a template and render it
// first, otherwise it will just return a regular query.
prep := func(wrapped interface{}) (uint64, *structs.PreparedQuery, error) {
wrapper := wrapped.(*queryWrapper)
if prepared_query.IsTemplate(wrapper.PreparedQuery) {
render, err := wrapper.ct.Render(queryIDOrName, source)
if err != nil {
return idx, nil, err
}
return idx, render, nil
}
return idx, wrapper.PreparedQuery, nil
}
// Next, look for an exact name match. This is the common case for static
// prepared queries, and could also apply to templates.
{
wrapped, err := tx.First("prepared-queries", "name", queryIDOrName)
if err != nil {
return 0, nil, fmt.Errorf("failed prepared query lookup: %s", err)
}
if wrapped != nil {
return prep(wrapped)
}
}
// Next, look for the longest prefix match among the prepared query
// templates.
{
wrapped, err := tx.LongestPrefix("prepared-queries", "template_prefix", queryIDOrName)
if err != nil {
return 0, nil, fmt.Errorf("failed prepared query lookup: %s", err)
}
if wrapped != nil {
return prep(wrapped)
}
}
return idx, nil, nil
}
// PreparedQueryList returns all the prepared queries.
func (s *Store) PreparedQueryList(ws memdb.WatchSet) (uint64, structs.PreparedQueries, error) {
tx := s.db.Txn(false)
defer tx.Abort()
// Get the table index.
idx := maxIndexTxn(tx, "prepared-queries")
// Query all of the prepared queries in the state store.
queries, err := tx.Get("prepared-queries", "id")
if err != nil {
return 0, nil, fmt.Errorf("failed prepared query lookup: %s", err)
}
ws.Add(queries.WatchCh())
// Go over all of the queries and build the response.
var result structs.PreparedQueries
for wrapped := queries.Next(); wrapped != nil; wrapped = queries.Next() {
result = append(result, toPreparedQuery(wrapped))
}
return idx, result, nil
}