PageRenderTime 910ms CodeModel.GetById 14ms RepoModel.GetById 0ms app.codeStats 0ms

/vendor/github.com/hashicorp/consul/consul/state/prepared_query.go

https://github.com/backstage/backstage
Go | 353 lines | 235 code | 46 blank | 72 comment | 101 complexity | 01f076e089e238f9eb9a75740c5d3eef MD5 | raw file
Possible License(s): Apache-2.0, MIT, BSD-3-Clause, MPL-2.0-no-copyleft-exception
  1. package state
  2. import (
  3. "fmt"
  4. "regexp"
  5. "github.com/hashicorp/consul/consul/prepared_query"
  6. "github.com/hashicorp/consul/consul/structs"
  7. "github.com/hashicorp/go-memdb"
  8. )
  9. // validUUID is used to check if a given string looks like a UUID
  10. var validUUID = regexp.MustCompile(`(?i)^[\da-f]{8}-[\da-f]{4}-[\da-f]{4}-[\da-f]{4}-[\da-f]{12}$`)
  11. // isUUID returns true if the given string is a valid UUID.
  12. func isUUID(str string) bool {
  13. return validUUID.MatchString(str)
  14. }
  15. // queryWrapper is an internal structure that is used to store a query alongside
  16. // its compiled template, which can be nil.
  17. type queryWrapper struct {
  18. // We embed the PreparedQuery structure so that the UUID field indexer
  19. // can see the ID directly.
  20. *structs.PreparedQuery
  21. // ct is the compiled template, or nil if the query isn't a template. The
  22. // state store manages this and keeps it up to date every time the query
  23. // changes.
  24. ct *prepared_query.CompiledTemplate
  25. }
  26. // toPreparedQuery unwraps the internal form of a prepared query and returns
  27. // the regular struct.
  28. func toPreparedQuery(wrapped interface{}) *structs.PreparedQuery {
  29. if wrapped == nil {
  30. return nil
  31. }
  32. return wrapped.(*queryWrapper).PreparedQuery
  33. }
  34. // PreparedQueries is used to pull all the prepared queries from the snapshot.
  35. func (s *StateSnapshot) PreparedQueries() (structs.PreparedQueries, error) {
  36. queries, err := s.tx.Get("prepared-queries", "id")
  37. if err != nil {
  38. return nil, err
  39. }
  40. var ret structs.PreparedQueries
  41. for wrapped := queries.Next(); wrapped != nil; wrapped = queries.Next() {
  42. ret = append(ret, toPreparedQuery(wrapped))
  43. }
  44. return ret, nil
  45. }
  46. // PrepparedQuery is used when restoring from a snapshot. For general inserts,
  47. // use PreparedQuerySet.
  48. func (s *StateRestore) PreparedQuery(query *structs.PreparedQuery) error {
  49. // If this is a template, compile it, otherwise leave the compiled
  50. // template field nil.
  51. var ct *prepared_query.CompiledTemplate
  52. if prepared_query.IsTemplate(query) {
  53. var err error
  54. ct, err = prepared_query.Compile(query)
  55. if err != nil {
  56. return fmt.Errorf("failed compiling template: %s", err)
  57. }
  58. }
  59. // Insert the wrapped query.
  60. if err := s.tx.Insert("prepared-queries", &queryWrapper{query, ct}); err != nil {
  61. return fmt.Errorf("failed restoring prepared query: %s", err)
  62. }
  63. if err := indexUpdateMaxTxn(s.tx, query.ModifyIndex, "prepared-queries"); err != nil {
  64. return fmt.Errorf("failed updating index: %s", err)
  65. }
  66. s.watches.Arm("prepared-queries")
  67. return nil
  68. }
  69. // PreparedQuerySet is used to create or update a prepared query.
  70. func (s *StateStore) PreparedQuerySet(idx uint64, query *structs.PreparedQuery) error {
  71. tx := s.db.Txn(true)
  72. defer tx.Abort()
  73. if err := s.preparedQuerySetTxn(tx, idx, query); err != nil {
  74. return err
  75. }
  76. tx.Commit()
  77. return nil
  78. }
  79. // preparedQuerySetTxn is the inner method used to insert a prepared query with
  80. // the proper indexes into the state store.
  81. func (s *StateStore) preparedQuerySetTxn(tx *memdb.Txn, idx uint64, query *structs.PreparedQuery) error {
  82. // Check that the ID is set.
  83. if query.ID == "" {
  84. return ErrMissingQueryID
  85. }
  86. // Check for an existing query.
  87. wrapped, err := tx.First("prepared-queries", "id", query.ID)
  88. if err != nil {
  89. return fmt.Errorf("failed prepared query lookup: %s", err)
  90. }
  91. existing := toPreparedQuery(wrapped)
  92. // Set the indexes.
  93. if existing != nil {
  94. query.CreateIndex = existing.CreateIndex
  95. query.ModifyIndex = idx
  96. } else {
  97. query.CreateIndex = idx
  98. query.ModifyIndex = idx
  99. }
  100. // Verify that the query name doesn't already exist, or that we are
  101. // updating the same instance that has this name. If this is a template
  102. // and the name is empty then we make sure there's not an empty template
  103. // already registered.
  104. if query.Name != "" {
  105. wrapped, err := tx.First("prepared-queries", "name", query.Name)
  106. if err != nil {
  107. return fmt.Errorf("failed prepared query lookup: %s", err)
  108. }
  109. other := toPreparedQuery(wrapped)
  110. if other != nil && (existing == nil || existing.ID != other.ID) {
  111. return fmt.Errorf("name '%s' aliases an existing query name", query.Name)
  112. }
  113. } else if prepared_query.IsTemplate(query) {
  114. wrapped, err := tx.First("prepared-queries", "template", query.Name)
  115. if err != nil {
  116. return fmt.Errorf("failed prepared query lookup: %s", err)
  117. }
  118. other := toPreparedQuery(wrapped)
  119. if other != nil && (existing == nil || existing.ID != other.ID) {
  120. return fmt.Errorf("a query template with an empty name already exists")
  121. }
  122. }
  123. // Verify that the name doesn't alias any existing ID. We allow queries
  124. // to be looked up by ID *or* name so we don't want anyone to try to
  125. // register a query with a name equal to some other query's ID in an
  126. // attempt to hijack it. We also look up by ID *then* name in order to
  127. // prevent this, but it seems prudent to prevent these types of rogue
  128. // queries from ever making it into the state store. Note that we have
  129. // to see if the name looks like a UUID before checking since the UUID
  130. // index will complain if we look up something that's not formatted
  131. // like one.
  132. if isUUID(query.Name) {
  133. wrapped, err := tx.First("prepared-queries", "id", query.Name)
  134. if err != nil {
  135. return fmt.Errorf("failed prepared query lookup: %s", err)
  136. }
  137. if wrapped != nil {
  138. return fmt.Errorf("name '%s' aliases an existing query ID", query.Name)
  139. }
  140. }
  141. // Verify that the session exists.
  142. if query.Session != "" {
  143. sess, err := tx.First("sessions", "id", query.Session)
  144. if err != nil {
  145. return fmt.Errorf("failed session lookup: %s", err)
  146. }
  147. if sess == nil {
  148. return fmt.Errorf("invalid session %#v", query.Session)
  149. }
  150. }
  151. // We do not verify the service here, nor the token, if any. These are
  152. // checked at execute time and not doing integrity checking on them
  153. // helps avoid bootstrapping chicken and egg problems.
  154. // If this is a template, compile it, otherwise leave the compiled
  155. // template field nil.
  156. var ct *prepared_query.CompiledTemplate
  157. if prepared_query.IsTemplate(query) {
  158. var err error
  159. ct, err = prepared_query.Compile(query)
  160. if err != nil {
  161. return fmt.Errorf("failed compiling template: %s", err)
  162. }
  163. }
  164. // Insert the wrapped query.
  165. if err := tx.Insert("prepared-queries", &queryWrapper{query, ct}); err != nil {
  166. return fmt.Errorf("failed inserting prepared query: %s", err)
  167. }
  168. if err := tx.Insert("index", &IndexEntry{"prepared-queries", idx}); err != nil {
  169. return fmt.Errorf("failed updating index: %s", err)
  170. }
  171. tx.Defer(func() { s.tableWatches["prepared-queries"].Notify() })
  172. return nil
  173. }
  174. // PreparedQueryDelete deletes the given query by ID.
  175. func (s *StateStore) PreparedQueryDelete(idx uint64, queryID string) error {
  176. tx := s.db.Txn(true)
  177. defer tx.Abort()
  178. watches := NewDumbWatchManager(s.tableWatches)
  179. if err := s.preparedQueryDeleteTxn(tx, idx, watches, queryID); err != nil {
  180. return fmt.Errorf("failed prepared query delete: %s", err)
  181. }
  182. tx.Defer(func() { watches.Notify() })
  183. tx.Commit()
  184. return nil
  185. }
  186. // preparedQueryDeleteTxn is the inner method used to delete a prepared query
  187. // with the proper indexes into the state store.
  188. func (s *StateStore) preparedQueryDeleteTxn(tx *memdb.Txn, idx uint64, watches *DumbWatchManager,
  189. queryID string) error {
  190. // Pull the query.
  191. wrapped, err := tx.First("prepared-queries", "id", queryID)
  192. if err != nil {
  193. return fmt.Errorf("failed prepared query lookup: %s", err)
  194. }
  195. if wrapped == nil {
  196. return nil
  197. }
  198. // Delete the query and update the index.
  199. if err := tx.Delete("prepared-queries", wrapped); err != nil {
  200. return fmt.Errorf("failed prepared query delete: %s", err)
  201. }
  202. if err := tx.Insert("index", &IndexEntry{"prepared-queries", idx}); err != nil {
  203. return fmt.Errorf("failed updating index: %s", err)
  204. }
  205. watches.Arm("prepared-queries")
  206. return nil
  207. }
  208. // PreparedQueryGet returns the given prepared query by ID.
  209. func (s *StateStore) PreparedQueryGet(queryID string) (uint64, *structs.PreparedQuery, error) {
  210. tx := s.db.Txn(false)
  211. defer tx.Abort()
  212. // Get the table index.
  213. idx := maxIndexTxn(tx, s.getWatchTables("PreparedQueryGet")...)
  214. // Look up the query by its ID.
  215. wrapped, err := tx.First("prepared-queries", "id", queryID)
  216. if err != nil {
  217. return 0, nil, fmt.Errorf("failed prepared query lookup: %s", err)
  218. }
  219. return idx, toPreparedQuery(wrapped), nil
  220. }
  221. // PreparedQueryResolve returns the given prepared query by looking up an ID or
  222. // Name. If the query was looked up by name and it's a template, then the
  223. // template will be rendered before it is returned.
  224. func (s *StateStore) PreparedQueryResolve(queryIDOrName string) (uint64, *structs.PreparedQuery, error) {
  225. tx := s.db.Txn(false)
  226. defer tx.Abort()
  227. // Get the table index.
  228. idx := maxIndexTxn(tx, s.getWatchTables("PreparedQueryResolve")...)
  229. // Explicitly ban an empty query. This will never match an ID and the
  230. // schema is set up so it will never match a query with an empty name,
  231. // but we check it here to be explicit about it (we'd never want to
  232. // return the results from the first query w/o a name).
  233. if queryIDOrName == "" {
  234. return 0, nil, ErrMissingQueryID
  235. }
  236. // Try first by ID if it looks like they gave us an ID. We check the
  237. // format before trying this because the UUID index will complain if
  238. // we look up something that's not formatted like one.
  239. if isUUID(queryIDOrName) {
  240. wrapped, err := tx.First("prepared-queries", "id", queryIDOrName)
  241. if err != nil {
  242. return 0, nil, fmt.Errorf("failed prepared query lookup: %s", err)
  243. }
  244. if wrapped != nil {
  245. query := toPreparedQuery(wrapped)
  246. if prepared_query.IsTemplate(query) {
  247. return idx, nil, fmt.Errorf("prepared query templates can only be resolved up by name, not by ID")
  248. }
  249. return idx, query, nil
  250. }
  251. }
  252. // prep will check to see if the query is a template and render it
  253. // first, otherwise it will just return a regular query.
  254. prep := func(wrapped interface{}) (uint64, *structs.PreparedQuery, error) {
  255. wrapper := wrapped.(*queryWrapper)
  256. if prepared_query.IsTemplate(wrapper.PreparedQuery) {
  257. render, err := wrapper.ct.Render(queryIDOrName)
  258. if err != nil {
  259. return idx, nil, err
  260. }
  261. return idx, render, nil
  262. } else {
  263. return idx, wrapper.PreparedQuery, nil
  264. }
  265. }
  266. // Next, look for an exact name match. This is the common case for static
  267. // prepared queries, and could also apply to templates.
  268. {
  269. wrapped, err := tx.First("prepared-queries", "name", queryIDOrName)
  270. if err != nil {
  271. return 0, nil, fmt.Errorf("failed prepared query lookup: %s", err)
  272. }
  273. if wrapped != nil {
  274. return prep(wrapped)
  275. }
  276. }
  277. // Next, look for the longest prefix match among the prepared query
  278. // templates.
  279. {
  280. wrapped, err := tx.LongestPrefix("prepared-queries", "template_prefix", queryIDOrName)
  281. if err != nil {
  282. return 0, nil, fmt.Errorf("failed prepared query lookup: %s", err)
  283. }
  284. if wrapped != nil {
  285. return prep(wrapped)
  286. }
  287. }
  288. return idx, nil, nil
  289. }
  290. // PreparedQueryList returns all the prepared queries.
  291. func (s *StateStore) PreparedQueryList() (uint64, structs.PreparedQueries, error) {
  292. tx := s.db.Txn(false)
  293. defer tx.Abort()
  294. // Get the table index.
  295. idx := maxIndexTxn(tx, s.getWatchTables("PreparedQueryList")...)
  296. // Query all of the prepared queries in the state store.
  297. queries, err := tx.Get("prepared-queries", "id")
  298. if err != nil {
  299. return 0, nil, fmt.Errorf("failed prepared query lookup: %s", err)
  300. }
  301. // Go over all of the queries and build the response.
  302. var result structs.PreparedQueries
  303. for wrapped := queries.Next(); wrapped != nil; wrapped = queries.Next() {
  304. result = append(result, toPreparedQuery(wrapped))
  305. }
  306. return idx, result, nil
  307. }