dataStore.go 41 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460
  1. /*
  2. * Copyright (c) 2015, Psiphon Inc.
  3. * All rights reserved.
  4. *
  5. * This program is free software: you can redistribute it and/or modify
  6. * it under the terms of the GNU General Public License as published by
  7. * the Free Software Foundation, either version 3 of the License, or
  8. * (at your option) any later version.
  9. *
  10. * This program is distributed in the hope that it will be useful,
  11. * but WITHOUT ANY WARRANTY; without even the implied warranty of
  12. * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  13. * GNU General Public License for more details.
  14. *
  15. * You should have received a copy of the GNU General Public License
  16. * along with this program. If not, see <http://www.gnu.org/licenses/>.
  17. *
  18. */
  19. package psiphon
  20. import (
  21. "bytes"
  22. "encoding/json"
  23. "errors"
  24. "fmt"
  25. "math/rand"
  26. "os"
  27. "path/filepath"
  28. "sync"
  29. "time"
  30. "github.com/Psiphon-Inc/bolt"
  31. "github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common"
  32. "github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/parameters"
  33. "github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/protocol"
  34. )
  35. // The BoltDB dataStore implementation is an alternative to the sqlite3-based
  36. // implementation in dataStore.go. Both implementations have the same interface.
  37. //
  38. // BoltDB is pure Go, and is intended to be used in cases where we have trouble
  39. // building sqlite3/CGO (e.g., currently go mobile due to
  40. // https://github.com/mattn/go-sqlite3/issues/201), and perhaps ultimately as
  41. // the primary dataStore implementation.
  42. //
  43. type dataStore struct {
  44. init sync.Once
  45. db *bolt.DB
  46. }
  47. const (
  48. serverEntriesBucket = "serverEntries"
  49. rankedServerEntriesBucket = "rankedServerEntries"
  50. rankedServerEntriesKey = "rankedServerEntries"
  51. splitTunnelRouteETagsBucket = "splitTunnelRouteETags"
  52. splitTunnelRouteDataBucket = "splitTunnelRouteData"
  53. urlETagsBucket = "urlETags"
  54. keyValueBucket = "keyValues"
  55. tunnelStatsBucket = "tunnelStats"
  56. remoteServerListStatsBucket = "remoteServerListStats"
  57. slokBucket = "SLOKs"
  58. tacticsBucket = "tactics"
  59. speedTestSamplesBucket = "speedTestSamples"
  60. rankedServerEntryCount = 100
  61. )
  62. const (
  63. DATA_STORE_FILENAME = "psiphon.boltdb"
  64. LEGACY_DATA_STORE_FILENAME = "psiphon.db"
  65. DATA_STORE_LAST_CONNECTED_KEY = "lastConnected"
  66. DATA_STORE_LAST_SERVER_ENTRY_FILTER_KEY = "lastServerEntryFilter"
  67. PERSISTENT_STAT_TYPE_REMOTE_SERVER_LIST = remoteServerListStatsBucket
  68. )
  69. var singleton dataStore
  70. // InitDataStore initializes the singleton instance of dataStore. This
  71. // function uses a sync.Once and is safe for use by concurrent goroutines.
  72. // The underlying sql.DB connection pool is also safe.
  73. //
  74. // Note: the sync.Once was more useful when initDataStore was private and
  75. // called on-demand by the public functions below. Now we require an explicit
  76. // InitDataStore() call with the filename passed in. The on-demand calls
  77. // have been replaced by checkInitDataStore() to assert that Init was called.
  78. func InitDataStore(config *Config) (err error) {
  79. singleton.init.Do(func() {
  80. // Need to gather the list of migratable server entries before
  81. // initializing the boltdb store (as prepareMigrationEntries
  82. // checks for the existence of the bolt db file)
  83. migratableServerEntries := prepareMigrationEntries(config)
  84. filename := filepath.Join(config.DataStoreDirectory, DATA_STORE_FILENAME)
  85. var db *bolt.DB
  86. for retry := 0; retry < 3; retry++ {
  87. if retry > 0 {
  88. NoticeAlert("InitDataStore retry: %d", retry)
  89. }
  90. db, err = bolt.Open(filename, 0600, &bolt.Options{Timeout: 1 * time.Second})
  91. // The datastore file may be corrupt, so attempt to delete and try again
  92. if err != nil {
  93. NoticeAlert("bolt.Open error: %s", err)
  94. os.Remove(filename)
  95. continue
  96. }
  97. // Run consistency checks on datastore and emit errors for diagnostics purposes
  98. // We assume this will complete quickly for typical size Psiphon datastores.
  99. err = db.View(func(tx *bolt.Tx) error {
  100. return tx.SynchronousCheck()
  101. })
  102. // The datastore file may be corrupt, so attempt to delete and try again
  103. if err != nil {
  104. NoticeAlert("bolt.SynchronousCheck error: %s", err)
  105. db.Close()
  106. os.Remove(filename)
  107. continue
  108. }
  109. break
  110. }
  111. if err != nil {
  112. // Note: intending to set the err return value for InitDataStore
  113. err = fmt.Errorf("initDataStore failed to open database: %s", err)
  114. return
  115. }
  116. err = db.Update(func(tx *bolt.Tx) error {
  117. requiredBuckets := []string{
  118. serverEntriesBucket,
  119. rankedServerEntriesBucket,
  120. splitTunnelRouteETagsBucket,
  121. splitTunnelRouteDataBucket,
  122. urlETagsBucket,
  123. keyValueBucket,
  124. tunnelStatsBucket,
  125. remoteServerListStatsBucket,
  126. slokBucket,
  127. tacticsBucket,
  128. speedTestSamplesBucket,
  129. }
  130. for _, bucket := range requiredBuckets {
  131. _, err := tx.CreateBucketIfNotExists([]byte(bucket))
  132. if err != nil {
  133. return err
  134. }
  135. }
  136. return nil
  137. })
  138. if err != nil {
  139. err = fmt.Errorf("initDataStore failed to create buckets: %s", err)
  140. return
  141. }
  142. // Cleanup obsolete tunnel (session) stats bucket, if one still exists
  143. err = db.Update(func(tx *bolt.Tx) error {
  144. tunnelStatsBucket := []byte("tunnelStats")
  145. if tx.Bucket(tunnelStatsBucket) != nil {
  146. err := tx.DeleteBucket(tunnelStatsBucket)
  147. if err != nil {
  148. NoticeAlert("DeleteBucket %s error: %s", tunnelStatsBucket, err)
  149. // Continue, since this is not fatal
  150. }
  151. }
  152. return nil
  153. })
  154. if err != nil {
  155. err = fmt.Errorf("initDataStore failed to create buckets: %s", err)
  156. return
  157. }
  158. singleton.db = db
  159. // The migrateServerEntries function requires the data store is
  160. // initialized prior to execution so that migrated entries can be stored
  161. if len(migratableServerEntries) > 0 {
  162. migrateEntries(
  163. config, migratableServerEntries, filepath.Join(config.DataStoreDirectory, LEGACY_DATA_STORE_FILENAME))
  164. }
  165. resetAllPersistentStatsToUnreported()
  166. })
  167. return err
  168. }
  169. func checkInitDataStore() {
  170. if singleton.db == nil {
  171. panic("checkInitDataStore: datastore not initialized")
  172. }
  173. }
  174. // StoreServerEntry adds the server entry to the data store.
  175. // A newly stored (or re-stored) server entry is assigned the next-to-top
  176. // rank for iteration order (the previous top ranked entry is promoted). The
  177. // purpose of inserting at next-to-top is to keep the last selected server
  178. // as the top ranked server.
  179. //
  180. // When a server entry already exists for a given server, it will be
  181. // replaced only if replaceIfExists is set or if the the ConfigurationVersion
  182. // field of the new entry is strictly higher than the existing entry.
  183. //
  184. // If the server entry data is malformed, an alert notice is issued and
  185. // the entry is skipped; no error is returned.
  186. func StoreServerEntry(serverEntry *protocol.ServerEntry, replaceIfExists bool) error {
  187. checkInitDataStore()
  188. // Server entries should already be validated before this point,
  189. // so instead of skipping we fail with an error.
  190. err := protocol.ValidateServerEntry(serverEntry)
  191. if err != nil {
  192. return common.ContextError(
  193. fmt.Errorf("invalid server entry: %s", err))
  194. }
  195. // BoltDB implementation note:
  196. // For simplicity, we don't maintain indexes on server entry
  197. // region or supported protocols. Instead, we perform full-bucket
  198. // scans with a filter. With a small enough database (thousands or
  199. // even tens of thousand of server entries) and common enough
  200. // values (e.g., many servers support all protocols), performance
  201. // is expected to be acceptable.
  202. err = singleton.db.Update(func(tx *bolt.Tx) error {
  203. serverEntries := tx.Bucket([]byte(serverEntriesBucket))
  204. // Check not only that the entry exists, but is valid. This
  205. // will replace in the rare case where the data is corrupt.
  206. existingConfigurationVersion := -1
  207. existingData := serverEntries.Get([]byte(serverEntry.IpAddress))
  208. if existingData != nil {
  209. var existingServerEntry *protocol.ServerEntry
  210. err := json.Unmarshal(existingData, &existingServerEntry)
  211. if err == nil {
  212. existingConfigurationVersion = existingServerEntry.ConfigurationVersion
  213. }
  214. }
  215. exists := existingConfigurationVersion > -1
  216. newer := exists && existingConfigurationVersion < serverEntry.ConfigurationVersion
  217. update := !exists || replaceIfExists || newer
  218. if !update {
  219. // Disabling this notice, for now, as it generates too much noise
  220. // in diagnostics with clients that always submit embedded servers
  221. // to the core on each run.
  222. // NoticeInfo("ignored update for server %s", serverEntry.IpAddress)
  223. return nil
  224. }
  225. data, err := json.Marshal(serverEntry)
  226. if err != nil {
  227. return common.ContextError(err)
  228. }
  229. err = serverEntries.Put([]byte(serverEntry.IpAddress), data)
  230. if err != nil {
  231. return common.ContextError(err)
  232. }
  233. err = insertRankedServerEntry(tx, serverEntry.IpAddress, 1)
  234. if err != nil {
  235. return common.ContextError(err)
  236. }
  237. NoticeInfo("updated server %s", serverEntry.IpAddress)
  238. return nil
  239. })
  240. if err != nil {
  241. return common.ContextError(err)
  242. }
  243. return nil
  244. }
  245. // StoreServerEntries stores a list of server entries.
  246. // There is an independent transaction for each entry insert/update.
  247. func StoreServerEntries(
  248. config *Config,
  249. serverEntries []*protocol.ServerEntry,
  250. replaceIfExists bool) error {
  251. checkInitDataStore()
  252. for _, serverEntry := range serverEntries {
  253. err := StoreServerEntry(serverEntry, replaceIfExists)
  254. if err != nil {
  255. return common.ContextError(err)
  256. }
  257. }
  258. // Since there has possibly been a significant change in the server entries,
  259. // take this opportunity to update the available egress regions.
  260. ReportAvailableRegions(config)
  261. return nil
  262. }
  263. // StreamingStoreServerEntries stores a list of server entries.
  264. // There is an independent transaction for each entry insert/update.
  265. func StreamingStoreServerEntries(
  266. config *Config,
  267. serverEntries *protocol.StreamingServerEntryDecoder,
  268. replaceIfExists bool) error {
  269. checkInitDataStore()
  270. // Note: both StreamingServerEntryDecoder.Next and StoreServerEntry
  271. // allocate temporary memory buffers for hex/JSON decoding/encoding,
  272. // so this isn't true constant-memory streaming (it depends on garbage
  273. // collection).
  274. for {
  275. serverEntry, err := serverEntries.Next()
  276. if err != nil {
  277. return common.ContextError(err)
  278. }
  279. if serverEntry == nil {
  280. // No more server entries
  281. break
  282. }
  283. err = StoreServerEntry(serverEntry, replaceIfExists)
  284. if err != nil {
  285. return common.ContextError(err)
  286. }
  287. }
  288. // Since there has possibly been a significant change in the server entries,
  289. // take this opportunity to update the available egress regions.
  290. ReportAvailableRegions(config)
  291. return nil
  292. }
  293. // PromoteServerEntry assigns the top rank (one more than current
  294. // max rank) to the specified server entry. Server candidates are
  295. // iterated in decending rank order, so this server entry will be
  296. // the first candidate in a subsequent tunnel establishment.
  297. func PromoteServerEntry(config *Config, ipAddress string) error {
  298. checkInitDataStore()
  299. err := singleton.db.Update(func(tx *bolt.Tx) error {
  300. // Ensure the corresponding entry exists before
  301. // inserting into rank.
  302. bucket := tx.Bucket([]byte(serverEntriesBucket))
  303. data := bucket.Get([]byte(ipAddress))
  304. if data == nil {
  305. NoticeAlert(
  306. "PromoteServerEntry: ignoring unknown server entry: %s",
  307. ipAddress)
  308. return nil
  309. }
  310. err := insertRankedServerEntry(tx, ipAddress, 0)
  311. if err != nil {
  312. return err
  313. }
  314. // Store the current server entry filter (e.g, region, etc.) that
  315. // was in use when the entry was promoted. This is used to detect
  316. // when the top ranked server entry was promoted under a different
  317. // filter.
  318. currentFilter, err := makeServerEntryFilterValue(config)
  319. if err != nil {
  320. return err
  321. }
  322. bucket = tx.Bucket([]byte(keyValueBucket))
  323. return bucket.Put([]byte(DATA_STORE_LAST_SERVER_ENTRY_FILTER_KEY), currentFilter)
  324. })
  325. if err != nil {
  326. return common.ContextError(err)
  327. }
  328. return nil
  329. }
  330. func makeServerEntryFilterValue(config *Config) ([]byte, error) {
  331. // Currently, only a change of EgressRegion will "break" server affinity.
  332. // If the tunnel protocol filter changes, any existing affinity server
  333. // either passes the new filter, or it will be skipped anyway.
  334. return []byte(config.EgressRegion), nil
  335. }
  336. func hasServerEntryFilterChanged(config *Config) (bool, error) {
  337. currentFilter, err := makeServerEntryFilterValue(config)
  338. if err != nil {
  339. return false, common.ContextError(err)
  340. }
  341. changed := false
  342. err = singleton.db.View(func(tx *bolt.Tx) error {
  343. // previousFilter will be nil not found (not previously
  344. // set) which will never match any current filter.
  345. bucket := tx.Bucket([]byte(keyValueBucket))
  346. previousFilter := bucket.Get([]byte(DATA_STORE_LAST_SERVER_ENTRY_FILTER_KEY))
  347. if bytes.Compare(previousFilter, currentFilter) != 0 {
  348. changed = true
  349. }
  350. return nil
  351. })
  352. if err != nil {
  353. return false, common.ContextError(err)
  354. }
  355. return changed, nil
  356. }
  357. func getRankedServerEntries(tx *bolt.Tx) ([]string, error) {
  358. bucket := tx.Bucket([]byte(rankedServerEntriesBucket))
  359. data := bucket.Get([]byte(rankedServerEntriesKey))
  360. if data == nil {
  361. return []string{}, nil
  362. }
  363. rankedServerEntries := make([]string, 0)
  364. err := json.Unmarshal(data, &rankedServerEntries)
  365. if err != nil {
  366. return nil, common.ContextError(err)
  367. }
  368. return rankedServerEntries, nil
  369. }
  370. func setRankedServerEntries(tx *bolt.Tx, rankedServerEntries []string) error {
  371. data, err := json.Marshal(rankedServerEntries)
  372. if err != nil {
  373. return common.ContextError(err)
  374. }
  375. bucket := tx.Bucket([]byte(rankedServerEntriesBucket))
  376. err = bucket.Put([]byte(rankedServerEntriesKey), data)
  377. if err != nil {
  378. return common.ContextError(err)
  379. }
  380. return nil
  381. }
  382. func insertRankedServerEntry(tx *bolt.Tx, serverEntryId string, position int) error {
  383. rankedServerEntries, err := getRankedServerEntries(tx)
  384. if err != nil {
  385. return common.ContextError(err)
  386. }
  387. // BoltDB implementation note:
  388. // For simplicity, we store the ranked server ids in an array serialized to
  389. // a single key value. To ensure this value doesn't grow without bound,
  390. // it's capped at rankedServerEntryCount. For now, this cap should be large
  391. // enough to meet the shuffleHeadLength = config.TunnelPoolSize criteria, for
  392. // any reasonable configuration of config.TunnelPoolSize.
  393. // Using: https://github.com/golang/go/wiki/SliceTricks
  394. // When serverEntryId is already ranked, remove it first to avoid duplicates
  395. for i, rankedServerEntryId := range rankedServerEntries {
  396. if rankedServerEntryId == serverEntryId {
  397. rankedServerEntries = append(
  398. rankedServerEntries[:i], rankedServerEntries[i+1:]...)
  399. break
  400. }
  401. }
  402. // SliceTricks insert, with length cap enforced
  403. if len(rankedServerEntries) < rankedServerEntryCount {
  404. rankedServerEntries = append(rankedServerEntries, "")
  405. }
  406. if position >= len(rankedServerEntries) {
  407. position = len(rankedServerEntries) - 1
  408. }
  409. copy(rankedServerEntries[position+1:], rankedServerEntries[position:])
  410. rankedServerEntries[position] = serverEntryId
  411. err = setRankedServerEntries(tx, rankedServerEntries)
  412. if err != nil {
  413. return common.ContextError(err)
  414. }
  415. return nil
  416. }
  417. // ServerEntryIterator is used to iterate over
  418. // stored server entries in rank order.
  419. type ServerEntryIterator struct {
  420. config *Config
  421. shuffleHeadLength int
  422. serverEntryIds []string
  423. serverEntryIndex int
  424. isTacticsServerEntryIterator bool
  425. isTargetServerEntryIterator bool
  426. hasNextTargetServerEntry bool
  427. targetServerEntry *protocol.ServerEntry
  428. }
  429. // NewServerEntryIterator creates a new ServerEntryIterator.
  430. //
  431. // The boolean return value indicates whether to treat the first server(s)
  432. // as affinity servers or not. When the server entry selection filter changes
  433. // such as from a specific region to any region, or when there was no previous
  434. // filter/iterator, the the first server(s) are arbitrary and should not be
  435. // given affinity treatment.
  436. //
  437. // NewServerEntryIterator and any returned ServerEntryIterator are not
  438. // designed for concurrent use as not all related datastore operations are
  439. // performed in a single transaction.
  440. //
  441. func NewServerEntryIterator(config *Config) (bool, *ServerEntryIterator, error) {
  442. // When configured, this target server entry is the only candidate
  443. if config.TargetServerEntry != "" {
  444. return newTargetServerEntryIterator(config, false)
  445. }
  446. checkInitDataStore()
  447. filterChanged, err := hasServerEntryFilterChanged(config)
  448. if err != nil {
  449. return false, nil, common.ContextError(err)
  450. }
  451. applyServerAffinity := !filterChanged
  452. iterator := &ServerEntryIterator{
  453. config: config,
  454. shuffleHeadLength: config.TunnelPoolSize,
  455. }
  456. err = iterator.Reset()
  457. if err != nil {
  458. return false, nil, common.ContextError(err)
  459. }
  460. return applyServerAffinity, iterator, nil
  461. }
  462. func NewTacticsServerEntryIterator(config *Config) (*ServerEntryIterator, error) {
  463. // When configured, this target server entry is the only candidate
  464. if config.TargetServerEntry != "" {
  465. _, iterator, err := newTargetServerEntryIterator(config, true)
  466. return iterator, err
  467. }
  468. checkInitDataStore()
  469. iterator := &ServerEntryIterator{
  470. shuffleHeadLength: 0,
  471. isTacticsServerEntryIterator: true,
  472. }
  473. err := iterator.Reset()
  474. if err != nil {
  475. return nil, common.ContextError(err)
  476. }
  477. return iterator, nil
  478. }
  479. // newTargetServerEntryIterator is a helper for initializing the TargetServerEntry case
  480. func newTargetServerEntryIterator(config *Config, isTactics bool) (bool, *ServerEntryIterator, error) {
  481. serverEntry, err := protocol.DecodeServerEntry(
  482. config.TargetServerEntry, common.GetCurrentTimestamp(), protocol.SERVER_ENTRY_SOURCE_TARGET)
  483. if err != nil {
  484. return false, nil, common.ContextError(err)
  485. }
  486. if isTactics {
  487. if len(serverEntry.GetSupportedTacticsProtocols()) == 0 {
  488. return false, nil, common.ContextError(errors.New("TargetServerEntry does not support tactics protocols"))
  489. }
  490. } else {
  491. if config.EgressRegion != "" && serverEntry.Region != config.EgressRegion {
  492. return false, nil, common.ContextError(errors.New("TargetServerEntry does not support EgressRegion"))
  493. }
  494. limitTunnelProtocols := config.clientParameters.Get().TunnelProtocols(parameters.LimitTunnelProtocols)
  495. if len(limitTunnelProtocols) > 0 {
  496. // At the ServerEntryIterator level, only limitTunnelProtocols is applied;
  497. // impairedTunnelProtocols and excludeMeek are handled higher up.
  498. if len(serverEntry.GetSupportedProtocols(limitTunnelProtocols, nil, false)) == 0 {
  499. return false, nil, common.ContextError(errors.New("TargetServerEntry does not support LimitTunnelProtocols"))
  500. }
  501. }
  502. }
  503. iterator := &ServerEntryIterator{
  504. isTacticsServerEntryIterator: isTactics,
  505. isTargetServerEntryIterator: true,
  506. hasNextTargetServerEntry: true,
  507. targetServerEntry: serverEntry,
  508. }
  509. NoticeInfo("using TargetServerEntry: %s", serverEntry.IpAddress)
  510. return false, iterator, nil
  511. }
  512. // Reset a NewServerEntryIterator to the start of its cycle. The next
  513. // call to Next will return the first server entry.
  514. func (iterator *ServerEntryIterator) Reset() error {
  515. iterator.Close()
  516. if iterator.isTargetServerEntryIterator {
  517. iterator.hasNextTargetServerEntry = true
  518. return nil
  519. }
  520. // For diagnostics, it's useful to count the number of known server
  521. // entries that satisfy both the egress region and tunnel protocol
  522. // requirements. The tunnel protocol filter is not applied by the iterator
  523. // as protocol filtering, including impaire protocol and exclude-meek
  524. // logic, is all handled higher up.
  525. // TODO: for isTacticsServerEntryIterator, emit tactics candidate count.
  526. if !iterator.isTacticsServerEntryIterator {
  527. limitTunnelProtocols := iterator.config.clientParameters.Get().TunnelProtocols(
  528. parameters.LimitTunnelProtocols)
  529. count := CountServerEntries(iterator.config.EgressRegion, limitTunnelProtocols)
  530. NoticeCandidateServers(iterator.config.EgressRegion, limitTunnelProtocols, count)
  531. // LimitTunnelProtocols may have changed since the last ReportAvailableRegions,
  532. // and now there may be no servers with the required capabilities in the
  533. // selected region. ReportAvailableRegions will signal this to the client.
  534. if count == 0 {
  535. ReportAvailableRegions(iterator.config)
  536. }
  537. }
  538. // This query implements the Psiphon server candidate selection
  539. // algorithm: the first TunnelPoolSize server candidates are in rank
  540. // (priority) order, to favor previously successful servers; then the
  541. // remaining long tail is shuffled to raise up less recent candidates.
  542. // BoltDB implementation note:
  543. // We don't keep a transaction open for the duration of the iterator
  544. // because this would expose the following semantics to consumer code:
  545. //
  546. // Read-only transactions and read-write transactions ... generally
  547. // shouldn't be opened simultaneously in the same goroutine. This can
  548. // cause a deadlock as the read-write transaction needs to periodically
  549. // re-map the data file but it cannot do so while a read-only
  550. // transaction is open.
  551. // (https://github.com/boltdb/bolt)
  552. //
  553. // So the underlying serverEntriesBucket could change after the serverEntryIds
  554. // list is built.
  555. var serverEntryIds []string
  556. err := singleton.db.View(func(tx *bolt.Tx) error {
  557. var err error
  558. serverEntryIds, err = getRankedServerEntries(tx)
  559. if err != nil {
  560. return err
  561. }
  562. skipServerEntryIds := make(map[string]bool)
  563. for _, serverEntryId := range serverEntryIds {
  564. skipServerEntryIds[serverEntryId] = true
  565. }
  566. bucket := tx.Bucket([]byte(serverEntriesBucket))
  567. cursor := bucket.Cursor()
  568. for key, _ := cursor.Last(); key != nil; key, _ = cursor.Prev() {
  569. serverEntryId := string(key)
  570. if _, ok := skipServerEntryIds[serverEntryId]; ok {
  571. continue
  572. }
  573. serverEntryIds = append(serverEntryIds, serverEntryId)
  574. }
  575. return nil
  576. })
  577. if err != nil {
  578. return common.ContextError(err)
  579. }
  580. for i := len(serverEntryIds) - 1; i > iterator.shuffleHeadLength-1; i-- {
  581. j := rand.Intn(i+1-iterator.shuffleHeadLength) + iterator.shuffleHeadLength
  582. serverEntryIds[i], serverEntryIds[j] = serverEntryIds[j], serverEntryIds[i]
  583. }
  584. iterator.serverEntryIds = serverEntryIds
  585. iterator.serverEntryIndex = 0
  586. return nil
  587. }
  588. // Close cleans up resources associated with a ServerEntryIterator.
  589. func (iterator *ServerEntryIterator) Close() {
  590. iterator.serverEntryIds = nil
  591. iterator.serverEntryIndex = 0
  592. }
  593. // Next returns the next server entry, by rank, for a ServerEntryIterator.
  594. // Returns nil with no error when there is no next item.
  595. func (iterator *ServerEntryIterator) Next() (*protocol.ServerEntry, error) {
  596. var err error
  597. var serverEntry *protocol.ServerEntry
  598. defer func() {
  599. if err != nil {
  600. iterator.Close()
  601. }
  602. }()
  603. if iterator.isTargetServerEntryIterator {
  604. if iterator.hasNextTargetServerEntry {
  605. iterator.hasNextTargetServerEntry = false
  606. return MakeCompatibleServerEntry(iterator.targetServerEntry), nil
  607. }
  608. return nil, nil
  609. }
  610. // There are no region/protocol indexes for the server entries bucket.
  611. // Loop until we have the next server entry that matches the iterator
  612. // filter requirements.
  613. for {
  614. if iterator.serverEntryIndex >= len(iterator.serverEntryIds) {
  615. // There is no next item
  616. return nil, nil
  617. }
  618. serverEntryId := iterator.serverEntryIds[iterator.serverEntryIndex]
  619. iterator.serverEntryIndex += 1
  620. var data []byte
  621. err = singleton.db.View(func(tx *bolt.Tx) error {
  622. bucket := tx.Bucket([]byte(serverEntriesBucket))
  623. value := bucket.Get([]byte(serverEntryId))
  624. if value != nil {
  625. // Must make a copy as slice is only valid within transaction.
  626. data = make([]byte, len(value))
  627. copy(data, value)
  628. }
  629. return nil
  630. })
  631. if err != nil {
  632. return nil, common.ContextError(err)
  633. }
  634. if data == nil {
  635. // In case of data corruption or a bug causing this condition,
  636. // do not stop iterating.
  637. NoticeAlert("ServerEntryIterator.Next: unexpected missing server entry: %s", serverEntryId)
  638. continue
  639. }
  640. err = json.Unmarshal(data, &serverEntry)
  641. if err != nil {
  642. // In case of data corruption or a bug causing this condition,
  643. // do not stop iterating.
  644. NoticeAlert("ServerEntryIterator.Next: %s", common.ContextError(err))
  645. continue
  646. }
  647. // Check filter requirements
  648. if iterator.isTacticsServerEntryIterator {
  649. // Tactics doesn't filter by egress region.
  650. if len(serverEntry.GetSupportedTacticsProtocols()) > 0 {
  651. break
  652. }
  653. } else {
  654. if iterator.config.EgressRegion == "" ||
  655. serverEntry.Region == iterator.config.EgressRegion {
  656. break
  657. }
  658. }
  659. }
  660. return MakeCompatibleServerEntry(serverEntry), nil
  661. }
  662. // MakeCompatibleServerEntry provides backwards compatibility with old server entries
  663. // which have a single meekFrontingDomain and not a meekFrontingAddresses array.
  664. // By copying this one meekFrontingDomain into meekFrontingAddresses, this client effectively
  665. // uses that single value as legacy clients do.
  666. func MakeCompatibleServerEntry(serverEntry *protocol.ServerEntry) *protocol.ServerEntry {
  667. if len(serverEntry.MeekFrontingAddresses) == 0 && serverEntry.MeekFrontingDomain != "" {
  668. serverEntry.MeekFrontingAddresses =
  669. append(serverEntry.MeekFrontingAddresses, serverEntry.MeekFrontingDomain)
  670. }
  671. return serverEntry
  672. }
  673. func scanServerEntries(scanner func(*protocol.ServerEntry)) error {
  674. err := singleton.db.View(func(tx *bolt.Tx) error {
  675. bucket := tx.Bucket([]byte(serverEntriesBucket))
  676. cursor := bucket.Cursor()
  677. for key, value := cursor.First(); key != nil; key, value = cursor.Next() {
  678. serverEntry := new(protocol.ServerEntry)
  679. err := json.Unmarshal(value, serverEntry)
  680. if err != nil {
  681. // In case of data corruption or a bug causing this condition,
  682. // do not stop iterating.
  683. NoticeAlert("scanServerEntries: %s", common.ContextError(err))
  684. continue
  685. }
  686. scanner(serverEntry)
  687. }
  688. return nil
  689. })
  690. if err != nil {
  691. return common.ContextError(err)
  692. }
  693. return nil
  694. }
  695. // CountServerEntries returns a count of stored servers for the
  696. // specified region and tunnel protocols.
  697. func CountServerEntries(region string, tunnelProtocols []string) int {
  698. checkInitDataStore()
  699. count := 0
  700. err := scanServerEntries(func(serverEntry *protocol.ServerEntry) {
  701. if (region == "" || serverEntry.Region == region) &&
  702. (len(tunnelProtocols) == 0 ||
  703. // When CountServerEntries is called only limitTunnelProtocols is known;
  704. // impairedTunnelProtocols and excludeMeek may not apply.
  705. len(serverEntry.GetSupportedProtocols(tunnelProtocols, nil, false)) > 0) {
  706. count += 1
  707. }
  708. })
  709. if err != nil {
  710. NoticeAlert("CountServerEntries failed: %s", err)
  711. return 0
  712. }
  713. return count
  714. }
  715. // CountNonImpairedProtocols returns the number of distinct tunnel
  716. // protocols supported by stored server entries, excluding the
  717. // specified impaired protocols.
  718. func CountNonImpairedProtocols(
  719. region string,
  720. limitTunnelProtocols, impairedProtocols []string) int {
  721. checkInitDataStore()
  722. distinctProtocols := make(map[string]bool)
  723. err := scanServerEntries(func(serverEntry *protocol.ServerEntry) {
  724. if region == "" || serverEntry.Region == region {
  725. for _, protocol := range protocol.SupportedTunnelProtocols {
  726. if serverEntry.SupportsProtocol(protocol) {
  727. if len(limitTunnelProtocols) == 0 ||
  728. common.Contains(limitTunnelProtocols, protocol) {
  729. distinctProtocols[protocol] = true
  730. }
  731. }
  732. }
  733. }
  734. })
  735. for _, protocol := range impairedProtocols {
  736. delete(distinctProtocols, protocol)
  737. }
  738. if err != nil {
  739. NoticeAlert("CountNonImpairedProtocols failed: %s", err)
  740. return 0
  741. }
  742. return len(distinctProtocols)
  743. }
  744. // ReportAvailableRegions prints a notice with the available egress regions.
  745. func ReportAvailableRegions(config *Config) {
  746. checkInitDataStore()
  747. limitTunnelProtocols := config.clientParameters.Get().TunnelProtocols(
  748. parameters.LimitTunnelProtocols)
  749. regions := make(map[string]bool)
  750. err := scanServerEntries(func(serverEntry *protocol.ServerEntry) {
  751. if len(limitTunnelProtocols) == 0 ||
  752. // When ReportAvailableRegions is called only limitTunnelProtocols is known;
  753. // impairedTunnelProtocols and excludeMeek may not apply.
  754. len(serverEntry.GetSupportedProtocols(limitTunnelProtocols, nil, false)) > 0 {
  755. regions[serverEntry.Region] = true
  756. }
  757. })
  758. if err != nil {
  759. NoticeAlert("ReportAvailableRegions failed: %s", err)
  760. return
  761. }
  762. regionList := make([]string, 0, len(regions))
  763. for region := range regions {
  764. // Some server entries do not have a region, but it makes no sense to return
  765. // an empty string as an "available region".
  766. if region != "" {
  767. regionList = append(regionList, region)
  768. }
  769. }
  770. NoticeAvailableEgressRegions(regionList)
  771. }
  772. // GetServerEntryIpAddresses returns an array containing
  773. // all stored server IP addresses.
  774. func GetServerEntryIpAddresses() (ipAddresses []string, err error) {
  775. checkInitDataStore()
  776. ipAddresses = make([]string, 0)
  777. err = scanServerEntries(func(serverEntry *protocol.ServerEntry) {
  778. ipAddresses = append(ipAddresses, serverEntry.IpAddress)
  779. })
  780. if err != nil {
  781. return nil, common.ContextError(err)
  782. }
  783. return ipAddresses, nil
  784. }
  785. // SetSplitTunnelRoutes updates the cached routes data for
  786. // the given region. The associated etag is also stored and
  787. // used to make efficient web requests for updates to the data.
  788. func SetSplitTunnelRoutes(region, etag string, data []byte) error {
  789. checkInitDataStore()
  790. err := singleton.db.Update(func(tx *bolt.Tx) error {
  791. bucket := tx.Bucket([]byte(splitTunnelRouteETagsBucket))
  792. err := bucket.Put([]byte(region), []byte(etag))
  793. bucket = tx.Bucket([]byte(splitTunnelRouteDataBucket))
  794. err = bucket.Put([]byte(region), data)
  795. return err
  796. })
  797. if err != nil {
  798. return common.ContextError(err)
  799. }
  800. return nil
  801. }
  802. // GetSplitTunnelRoutesETag retrieves the etag for cached routes
  803. // data for the specified region. If not found, it returns an empty string value.
  804. func GetSplitTunnelRoutesETag(region string) (etag string, err error) {
  805. checkInitDataStore()
  806. err = singleton.db.View(func(tx *bolt.Tx) error {
  807. bucket := tx.Bucket([]byte(splitTunnelRouteETagsBucket))
  808. etag = string(bucket.Get([]byte(region)))
  809. return nil
  810. })
  811. if err != nil {
  812. return "", common.ContextError(err)
  813. }
  814. return etag, nil
  815. }
  816. // GetSplitTunnelRoutesData retrieves the cached routes data
  817. // for the specified region. If not found, it returns a nil value.
  818. func GetSplitTunnelRoutesData(region string) (data []byte, err error) {
  819. checkInitDataStore()
  820. err = singleton.db.View(func(tx *bolt.Tx) error {
  821. bucket := tx.Bucket([]byte(splitTunnelRouteDataBucket))
  822. value := bucket.Get([]byte(region))
  823. if value != nil {
  824. // Must make a copy as slice is only valid within transaction.
  825. data = make([]byte, len(value))
  826. copy(data, value)
  827. }
  828. return nil
  829. })
  830. if err != nil {
  831. return nil, common.ContextError(err)
  832. }
  833. return data, nil
  834. }
  835. // SetUrlETag stores an ETag for the specfied URL.
  836. // Note: input URL is treated as a string, and is not
  837. // encoded or decoded or otherwise canonicalized.
  838. func SetUrlETag(url, etag string) error {
  839. checkInitDataStore()
  840. err := singleton.db.Update(func(tx *bolt.Tx) error {
  841. bucket := tx.Bucket([]byte(urlETagsBucket))
  842. err := bucket.Put([]byte(url), []byte(etag))
  843. return err
  844. })
  845. if err != nil {
  846. return common.ContextError(err)
  847. }
  848. return nil
  849. }
  850. // GetUrlETag retrieves a previously stored an ETag for the
  851. // specfied URL. If not found, it returns an empty string value.
  852. func GetUrlETag(url string) (etag string, err error) {
  853. checkInitDataStore()
  854. err = singleton.db.View(func(tx *bolt.Tx) error {
  855. bucket := tx.Bucket([]byte(urlETagsBucket))
  856. etag = string(bucket.Get([]byte(url)))
  857. return nil
  858. })
  859. if err != nil {
  860. return "", common.ContextError(err)
  861. }
  862. return etag, nil
  863. }
  864. // SetKeyValue stores a key/value pair.
  865. func SetKeyValue(key, value string) error {
  866. checkInitDataStore()
  867. err := singleton.db.Update(func(tx *bolt.Tx) error {
  868. bucket := tx.Bucket([]byte(keyValueBucket))
  869. err := bucket.Put([]byte(key), []byte(value))
  870. return err
  871. })
  872. if err != nil {
  873. return common.ContextError(err)
  874. }
  875. return nil
  876. }
  877. // GetKeyValue retrieves the value for a given key. If not found,
  878. // it returns an empty string value.
  879. func GetKeyValue(key string) (value string, err error) {
  880. checkInitDataStore()
  881. err = singleton.db.View(func(tx *bolt.Tx) error {
  882. bucket := tx.Bucket([]byte(keyValueBucket))
  883. value = string(bucket.Get([]byte(key)))
  884. return nil
  885. })
  886. if err != nil {
  887. return "", common.ContextError(err)
  888. }
  889. return value, nil
  890. }
  891. // Persistent stat records in the persistentStatStateUnreported
  892. // state are available for take out.
  893. //
  894. // Records in the persistentStatStateReporting have been taken
  895. // out and are pending either deletion (for a successful request)
  896. // or change to StateUnreported (for a failed request).
  897. //
  898. // All persistent stat records are reverted to StateUnreported
  899. // when the datastore is initialized at start up.
  900. var persistentStatStateUnreported = []byte("0")
  901. var persistentStatStateReporting = []byte("1")
  902. var persistentStatTypes = []string{
  903. PERSISTENT_STAT_TYPE_REMOTE_SERVER_LIST,
  904. }
  905. // StorePersistentStat adds a new persistent stat record, which
  906. // is set to StateUnreported and is an immediate candidate for
  907. // reporting.
  908. //
  909. // The stat is a JSON byte array containing fields as
  910. // required by the Psiphon server API. It's assumed that the
  911. // JSON value contains enough unique information for the value to
  912. // function as a key in the key/value datastore. This assumption
  913. // is currently satisfied by the fields sessionId + tunnelNumber
  914. // for tunnel stats, and URL + ETag for remote server list stats.
  915. func StorePersistentStat(statType string, stat []byte) error {
  916. checkInitDataStore()
  917. if !common.Contains(persistentStatTypes, statType) {
  918. return common.ContextError(fmt.Errorf("invalid persistent stat type: %s", statType))
  919. }
  920. err := singleton.db.Update(func(tx *bolt.Tx) error {
  921. bucket := tx.Bucket([]byte(statType))
  922. err := bucket.Put(stat, persistentStatStateUnreported)
  923. return err
  924. })
  925. if err != nil {
  926. return common.ContextError(err)
  927. }
  928. return nil
  929. }
  930. // CountUnreportedPersistentStats returns the number of persistent
  931. // stat records in StateUnreported.
  932. func CountUnreportedPersistentStats() int {
  933. checkInitDataStore()
  934. unreported := 0
  935. err := singleton.db.View(func(tx *bolt.Tx) error {
  936. for _, statType := range persistentStatTypes {
  937. bucket := tx.Bucket([]byte(statType))
  938. cursor := bucket.Cursor()
  939. for key, value := cursor.First(); key != nil; key, value = cursor.Next() {
  940. if 0 == bytes.Compare(value, persistentStatStateUnreported) {
  941. unreported++
  942. break
  943. }
  944. }
  945. }
  946. return nil
  947. })
  948. if err != nil {
  949. NoticeAlert("CountUnreportedPersistentStats failed: %s", err)
  950. return 0
  951. }
  952. return unreported
  953. }
  954. // TakeOutUnreportedPersistentStats returns up to maxCount persistent
  955. // stats records that are in StateUnreported. The records are set to
  956. // StateReporting. If the records are successfully reported, clear them
  957. // with ClearReportedPersistentStats. If the records are not successfully
  958. // reported, restore them with PutBackUnreportedPersistentStats.
  959. func TakeOutUnreportedPersistentStats(maxCount int) (map[string][][]byte, error) {
  960. checkInitDataStore()
  961. stats := make(map[string][][]byte)
  962. err := singleton.db.Update(func(tx *bolt.Tx) error {
  963. count := 0
  964. for _, statType := range persistentStatTypes {
  965. bucket := tx.Bucket([]byte(statType))
  966. cursor := bucket.Cursor()
  967. for key, value := cursor.First(); key != nil; key, value = cursor.Next() {
  968. if count >= maxCount {
  969. break
  970. }
  971. // Perform a test JSON unmarshaling. In case of data corruption or a bug,
  972. // skip the record.
  973. var jsonData interface{}
  974. err := json.Unmarshal(key, &jsonData)
  975. if err != nil {
  976. NoticeAlert(
  977. "Invalid key in TakeOutUnreportedPersistentStats: %s: %s",
  978. string(key), err)
  979. continue
  980. }
  981. if 0 == bytes.Compare(value, persistentStatStateUnreported) {
  982. // Must make a copy as slice is only valid within transaction.
  983. data := make([]byte, len(key))
  984. copy(data, key)
  985. if stats[statType] == nil {
  986. stats[statType] = make([][]byte, 0)
  987. }
  988. stats[statType] = append(stats[statType], data)
  989. count += 1
  990. }
  991. }
  992. for _, key := range stats[statType] {
  993. err := bucket.Put(key, persistentStatStateReporting)
  994. if err != nil {
  995. return err
  996. }
  997. }
  998. }
  999. return nil
  1000. })
  1001. if err != nil {
  1002. return nil, common.ContextError(err)
  1003. }
  1004. return stats, nil
  1005. }
  1006. // PutBackUnreportedPersistentStats restores a list of persistent
  1007. // stat records to StateUnreported.
  1008. func PutBackUnreportedPersistentStats(stats map[string][][]byte) error {
  1009. checkInitDataStore()
  1010. err := singleton.db.Update(func(tx *bolt.Tx) error {
  1011. for _, statType := range persistentStatTypes {
  1012. bucket := tx.Bucket([]byte(statType))
  1013. for _, key := range stats[statType] {
  1014. err := bucket.Put(key, persistentStatStateUnreported)
  1015. if err != nil {
  1016. return err
  1017. }
  1018. }
  1019. }
  1020. return nil
  1021. })
  1022. if err != nil {
  1023. return common.ContextError(err)
  1024. }
  1025. return nil
  1026. }
  1027. // ClearReportedPersistentStats deletes a list of persistent
  1028. // stat records that were successfully reported.
  1029. func ClearReportedPersistentStats(stats map[string][][]byte) error {
  1030. checkInitDataStore()
  1031. err := singleton.db.Update(func(tx *bolt.Tx) error {
  1032. for _, statType := range persistentStatTypes {
  1033. bucket := tx.Bucket([]byte(statType))
  1034. for _, key := range stats[statType] {
  1035. err := bucket.Delete(key)
  1036. if err != nil {
  1037. return err
  1038. }
  1039. }
  1040. }
  1041. return nil
  1042. })
  1043. if err != nil {
  1044. return common.ContextError(err)
  1045. }
  1046. return nil
  1047. }
  1048. // resetAllPersistentStatsToUnreported sets all persistent stat
  1049. // records to StateUnreported. This reset is called when the
  1050. // datastore is initialized at start up, as we do not know if
  1051. // persistent records in StateReporting were reported or not.
  1052. func resetAllPersistentStatsToUnreported() error {
  1053. checkInitDataStore()
  1054. err := singleton.db.Update(func(tx *bolt.Tx) error {
  1055. for _, statType := range persistentStatTypes {
  1056. bucket := tx.Bucket([]byte(statType))
  1057. resetKeys := make([][]byte, 0)
  1058. cursor := bucket.Cursor()
  1059. for key, _ := cursor.First(); key != nil; key, _ = cursor.Next() {
  1060. resetKeys = append(resetKeys, key)
  1061. }
  1062. // TODO: data mutation is done outside cursor. Is this
  1063. // strictly necessary in this case? As is, this means
  1064. // all stats need to be loaded into memory at once.
  1065. // https://godoc.org/github.com/boltdb/bolt#Cursor
  1066. for _, key := range resetKeys {
  1067. err := bucket.Put(key, persistentStatStateUnreported)
  1068. if err != nil {
  1069. return err
  1070. }
  1071. }
  1072. }
  1073. return nil
  1074. })
  1075. if err != nil {
  1076. return common.ContextError(err)
  1077. }
  1078. return nil
  1079. }
  1080. // CountSLOKs returns the total number of SLOK records.
  1081. func CountSLOKs() int {
  1082. checkInitDataStore()
  1083. count := 0
  1084. err := singleton.db.View(func(tx *bolt.Tx) error {
  1085. bucket := tx.Bucket([]byte(slokBucket))
  1086. cursor := bucket.Cursor()
  1087. for key, _ := cursor.First(); key != nil; key, _ = cursor.Next() {
  1088. count++
  1089. }
  1090. return nil
  1091. })
  1092. if err != nil {
  1093. NoticeAlert("CountSLOKs failed: %s", err)
  1094. return 0
  1095. }
  1096. return count
  1097. }
  1098. // DeleteSLOKs deletes all SLOK records.
  1099. func DeleteSLOKs() error {
  1100. checkInitDataStore()
  1101. err := singleton.db.Update(func(tx *bolt.Tx) error {
  1102. bucket := tx.Bucket([]byte(slokBucket))
  1103. return bucket.ForEach(
  1104. func(id, _ []byte) error {
  1105. return bucket.Delete(id)
  1106. })
  1107. })
  1108. if err != nil {
  1109. return common.ContextError(err)
  1110. }
  1111. return nil
  1112. }
  1113. // SetSLOK stores a SLOK key, referenced by its ID. The bool
  1114. // return value indicates whether the SLOK was already stored.
  1115. func SetSLOK(id, key []byte) (bool, error) {
  1116. checkInitDataStore()
  1117. var duplicate bool
  1118. err := singleton.db.Update(func(tx *bolt.Tx) error {
  1119. bucket := tx.Bucket([]byte(slokBucket))
  1120. duplicate = bucket.Get(id) != nil
  1121. err := bucket.Put([]byte(id), []byte(key))
  1122. return err
  1123. })
  1124. if err != nil {
  1125. return false, common.ContextError(err)
  1126. }
  1127. return duplicate, nil
  1128. }
  1129. // GetSLOK returns a SLOK key for the specified ID. The return
  1130. // value is nil if the SLOK is not found.
  1131. func GetSLOK(id []byte) (key []byte, err error) {
  1132. checkInitDataStore()
  1133. err = singleton.db.View(func(tx *bolt.Tx) error {
  1134. bucket := tx.Bucket([]byte(slokBucket))
  1135. key = bucket.Get(id)
  1136. return nil
  1137. })
  1138. if err != nil {
  1139. return nil, common.ContextError(err)
  1140. }
  1141. return key, nil
  1142. }
  1143. // TacticsStorer implements tactics.Storer.
  1144. type TacticsStorer struct {
  1145. }
  1146. func (t *TacticsStorer) SetTacticsRecord(networkID string, record []byte) error {
  1147. return setBucketValue([]byte(tacticsBucket), []byte(networkID), record)
  1148. }
  1149. func (t *TacticsStorer) GetTacticsRecord(networkID string) ([]byte, error) {
  1150. return getBucketValue([]byte(tacticsBucket), []byte(networkID))
  1151. }
  1152. func (t *TacticsStorer) SetSpeedTestSamplesRecord(networkID string, record []byte) error {
  1153. return setBucketValue([]byte(speedTestSamplesBucket), []byte(networkID), record)
  1154. }
  1155. func (t *TacticsStorer) GetSpeedTestSamplesRecord(networkID string) ([]byte, error) {
  1156. return getBucketValue([]byte(speedTestSamplesBucket), []byte(networkID))
  1157. }
  1158. // GetTacticsStorer creates a TacticsStorer.
  1159. func GetTacticsStorer() *TacticsStorer {
  1160. return &TacticsStorer{}
  1161. }
  1162. func setBucketValue(bucket, key, value []byte) error {
  1163. checkInitDataStore()
  1164. err := singleton.db.Update(func(tx *bolt.Tx) error {
  1165. bucket := tx.Bucket(bucket)
  1166. err := bucket.Put(key, value)
  1167. return err
  1168. })
  1169. if err != nil {
  1170. return common.ContextError(err)
  1171. }
  1172. return nil
  1173. }
  1174. func getBucketValue(bucket, key []byte) (value []byte, err error) {
  1175. checkInitDataStore()
  1176. err = singleton.db.View(func(tx *bolt.Tx) error {
  1177. bucket := tx.Bucket(bucket)
  1178. value = bucket.Get(key)
  1179. return nil
  1180. })
  1181. if err != nil {
  1182. return nil, common.ContextError(err)
  1183. }
  1184. return value, nil
  1185. }