dataStore.go 48 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735
  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. "sync"
  24. "time"
  25. "github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common"
  26. "github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/errors"
  27. "github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/parameters"
  28. "github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/prng"
  29. "github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/protocol"
  30. )
  31. var (
  32. datastoreServerEntriesBucket = []byte("serverEntries")
  33. datastoreServerEntryTagsBucket = []byte("serverEntryTags")
  34. datastoreServerEntryTombstoneTagsBucket = []byte("serverEntryTombstoneTags")
  35. datastoreSplitTunnelRouteETagsBucket = []byte("splitTunnelRouteETags")
  36. datastoreSplitTunnelRouteDataBucket = []byte("splitTunnelRouteData")
  37. datastoreUrlETagsBucket = []byte("urlETags")
  38. datastoreKeyValueBucket = []byte("keyValues")
  39. datastoreRemoteServerListStatsBucket = []byte("remoteServerListStats")
  40. datastoreFailedTunnelStatsBucket = []byte("failedTunnelStats")
  41. datastoreSLOKsBucket = []byte("SLOKs")
  42. datastoreTacticsBucket = []byte("tactics")
  43. datastoreSpeedTestSamplesBucket = []byte("speedTestSamples")
  44. datastoreDialParametersBucket = []byte("dialParameters")
  45. datastoreLastConnectedKey = "lastConnected"
  46. datastoreLastServerEntryFilterKey = []byte("lastServerEntryFilter")
  47. datastoreAffinityServerEntryIDKey = []byte("affinityServerEntryID")
  48. datastorePersistentStatTypeRemoteServerList = string(datastoreRemoteServerListStatsBucket)
  49. datastorePersistentStatTypeFailedTunnel = string(datastoreFailedTunnelStatsBucket)
  50. datastoreServerEntryFetchGCThreshold = 20
  51. datastoreMutex sync.RWMutex
  52. activeDatastoreDB *datastoreDB
  53. )
  54. // OpenDataStore opens and initializes the singleton data store instance.
  55. func OpenDataStore(config *Config) error {
  56. datastoreMutex.Lock()
  57. existingDB := activeDatastoreDB
  58. if existingDB != nil {
  59. datastoreMutex.Unlock()
  60. return errors.TraceNew("db already open")
  61. }
  62. newDB, err := datastoreOpenDB(config.GetDataStoreDirectory())
  63. if err != nil {
  64. datastoreMutex.Unlock()
  65. return errors.Trace(err)
  66. }
  67. activeDatastoreDB = newDB
  68. datastoreMutex.Unlock()
  69. _ = resetAllPersistentStatsToUnreported()
  70. return nil
  71. }
  72. // CloseDataStore closes the singleton data store instance, if open.
  73. func CloseDataStore() {
  74. datastoreMutex.Lock()
  75. defer datastoreMutex.Unlock()
  76. if activeDatastoreDB == nil {
  77. return
  78. }
  79. err := activeDatastoreDB.close()
  80. if err != nil {
  81. NoticeWarning("failed to close database: %s", errors.Trace(err))
  82. }
  83. activeDatastoreDB = nil
  84. }
  85. func datastoreView(fn func(tx *datastoreTx) error) error {
  86. datastoreMutex.RLock()
  87. defer datastoreMutex.RUnlock()
  88. if activeDatastoreDB == nil {
  89. return errors.TraceNew("database not open")
  90. }
  91. err := activeDatastoreDB.view(fn)
  92. if err != nil {
  93. err = errors.Trace(err)
  94. }
  95. return err
  96. }
  97. func datastoreUpdate(fn func(tx *datastoreTx) error) error {
  98. datastoreMutex.RLock()
  99. defer datastoreMutex.RUnlock()
  100. if activeDatastoreDB == nil {
  101. return errors.TraceNew("database not open")
  102. }
  103. err := activeDatastoreDB.update(fn)
  104. if err != nil {
  105. err = errors.Trace(err)
  106. }
  107. return err
  108. }
  109. // StoreServerEntry adds the server entry to the data store.
  110. //
  111. // When a server entry already exists for a given server, it will be
  112. // replaced only if replaceIfExists is set or if the the ConfigurationVersion
  113. // field of the new entry is strictly higher than the existing entry.
  114. //
  115. // If the server entry data is malformed, an alert notice is issued and
  116. // the entry is skipped; no error is returned.
  117. func StoreServerEntry(serverEntryFields protocol.ServerEntryFields, replaceIfExists bool) error {
  118. // TODO: call serverEntryFields.VerifySignature. At this time, we do not do
  119. // this as not all server entries have an individual signature field. All
  120. // StoreServerEntry callers either call VerifySignature or obtain server
  121. // entries from a trusted source (embedded in a signed client, or in a signed
  122. // authenticated package).
  123. // Server entries should already be validated before this point,
  124. // so instead of skipping we fail with an error.
  125. err := protocol.ValidateServerEntryFields(serverEntryFields)
  126. if err != nil {
  127. return errors.Tracef("invalid server entry: %s", err)
  128. }
  129. // BoltDB implementation note:
  130. // For simplicity, we don't maintain indexes on server entry
  131. // region or supported protocols. Instead, we perform full-bucket
  132. // scans with a filter. With a small enough database (thousands or
  133. // even tens of thousand of server entries) and common enough
  134. // values (e.g., many servers support all protocols), performance
  135. // is expected to be acceptable.
  136. err = datastoreUpdate(func(tx *datastoreTx) error {
  137. serverEntries := tx.bucket(datastoreServerEntriesBucket)
  138. serverEntryTags := tx.bucket(datastoreServerEntryTagsBucket)
  139. serverEntryTombstoneTags := tx.bucket(datastoreServerEntryTombstoneTagsBucket)
  140. serverEntryID := []byte(serverEntryFields.GetIPAddress())
  141. // Check not only that the entry exists, but is valid. This
  142. // will replace in the rare case where the data is corrupt.
  143. existingConfigurationVersion := -1
  144. existingData := serverEntries.get(serverEntryID)
  145. if existingData != nil {
  146. var existingServerEntry *protocol.ServerEntry
  147. err := json.Unmarshal(existingData, &existingServerEntry)
  148. if err == nil {
  149. existingConfigurationVersion = existingServerEntry.ConfigurationVersion
  150. }
  151. }
  152. exists := existingConfigurationVersion > -1
  153. newer := exists && existingConfigurationVersion < serverEntryFields.GetConfigurationVersion()
  154. update := !exists || replaceIfExists || newer
  155. if !update {
  156. return nil
  157. }
  158. serverEntryTag := serverEntryFields.GetTag()
  159. // Generate a derived tag when the server entry has no tag.
  160. if serverEntryTag == "" {
  161. serverEntryTag = protocol.GenerateServerEntryTag(
  162. serverEntryFields.GetIPAddress(),
  163. serverEntryFields.GetWebServerSecret())
  164. serverEntryFields.SetTag(serverEntryTag)
  165. }
  166. serverEntryTagBytes := []byte(serverEntryTag)
  167. // Ignore the server entry if it was previously pruned and a tombstone is
  168. // set.
  169. //
  170. // This logic is enforced only for embedded server entries, as all other
  171. // sources are considered to be definitive and non-stale. These exceptions
  172. // intentionally allow the scenario where a server is temporarily deleted
  173. // and then restored; in this case, it's desired for pruned server entries
  174. // to be restored.
  175. if serverEntryFields.GetLocalSource() == protocol.SERVER_ENTRY_SOURCE_EMBEDDED {
  176. if serverEntryTombstoneTags.get(serverEntryTagBytes) != nil {
  177. return nil
  178. }
  179. }
  180. data, err := json.Marshal(serverEntryFields)
  181. if err != nil {
  182. return errors.Trace(err)
  183. }
  184. err = serverEntries.put(serverEntryID, data)
  185. if err != nil {
  186. return errors.Trace(err)
  187. }
  188. err = serverEntryTags.put(serverEntryTagBytes, serverEntryID)
  189. if err != nil {
  190. return errors.Trace(err)
  191. }
  192. NoticeInfo("updated server %s", serverEntryFields.GetDiagnosticID())
  193. return nil
  194. })
  195. if err != nil {
  196. return errors.Trace(err)
  197. }
  198. return nil
  199. }
  200. // StoreServerEntries stores a list of server entries.
  201. // There is an independent transaction for each entry insert/update.
  202. func StoreServerEntries(
  203. config *Config,
  204. serverEntries []protocol.ServerEntryFields,
  205. replaceIfExists bool) error {
  206. for _, serverEntryFields := range serverEntries {
  207. err := StoreServerEntry(serverEntryFields, replaceIfExists)
  208. if err != nil {
  209. return errors.Trace(err)
  210. }
  211. }
  212. return nil
  213. }
  214. // StreamingStoreServerEntries stores a list of server entries.
  215. // There is an independent transaction for each entry insert/update.
  216. func StreamingStoreServerEntries(
  217. config *Config,
  218. serverEntries *protocol.StreamingServerEntryDecoder,
  219. replaceIfExists bool) error {
  220. // Note: both StreamingServerEntryDecoder.Next and StoreServerEntry
  221. // allocate temporary memory buffers for hex/JSON decoding/encoding,
  222. // so this isn't true constant-memory streaming (it depends on garbage
  223. // collection).
  224. n := 0
  225. for {
  226. serverEntry, err := serverEntries.Next()
  227. if err != nil {
  228. return errors.Trace(err)
  229. }
  230. if serverEntry == nil {
  231. // No more server entries
  232. break
  233. }
  234. err = StoreServerEntry(serverEntry, replaceIfExists)
  235. if err != nil {
  236. return errors.Trace(err)
  237. }
  238. n += 1
  239. if n == datastoreServerEntryFetchGCThreshold {
  240. DoGarbageCollection()
  241. n = 0
  242. }
  243. }
  244. return nil
  245. }
  246. // PromoteServerEntry sets the server affinity server entry ID to the
  247. // specified server entry IP address.
  248. func PromoteServerEntry(config *Config, ipAddress string) error {
  249. err := datastoreUpdate(func(tx *datastoreTx) error {
  250. serverEntryID := []byte(ipAddress)
  251. // Ensure the corresponding server entry exists before
  252. // setting server affinity.
  253. bucket := tx.bucket(datastoreServerEntriesBucket)
  254. data := bucket.get(serverEntryID)
  255. if data == nil {
  256. NoticeWarning(
  257. "PromoteServerEntry: ignoring unknown server entry: %s",
  258. ipAddress)
  259. return nil
  260. }
  261. bucket = tx.bucket(datastoreKeyValueBucket)
  262. err := bucket.put(datastoreAffinityServerEntryIDKey, serverEntryID)
  263. if err != nil {
  264. return errors.Trace(err)
  265. }
  266. // Store the current server entry filter (e.g, region, etc.) that
  267. // was in use when the entry was promoted. This is used to detect
  268. // when the top ranked server entry was promoted under a different
  269. // filter.
  270. currentFilter, err := makeServerEntryFilterValue(config)
  271. if err != nil {
  272. return errors.Trace(err)
  273. }
  274. return bucket.put(datastoreLastServerEntryFilterKey, currentFilter)
  275. })
  276. if err != nil {
  277. return errors.Trace(err)
  278. }
  279. return nil
  280. }
  281. func makeServerEntryFilterValue(config *Config) ([]byte, error) {
  282. // Currently, only a change of EgressRegion will "break" server affinity.
  283. // If the tunnel protocol filter changes, any existing affinity server
  284. // either passes the new filter, or it will be skipped anyway.
  285. return []byte(config.EgressRegion), nil
  286. }
  287. func hasServerEntryFilterChanged(config *Config) (bool, error) {
  288. currentFilter, err := makeServerEntryFilterValue(config)
  289. if err != nil {
  290. return false, errors.Trace(err)
  291. }
  292. changed := false
  293. err = datastoreView(func(tx *datastoreTx) error {
  294. bucket := tx.bucket(datastoreKeyValueBucket)
  295. previousFilter := bucket.get(datastoreLastServerEntryFilterKey)
  296. // When not found, previousFilter will be nil; ensures this
  297. // results in "changed", even if currentFilter is len(0).
  298. if previousFilter == nil ||
  299. !bytes.Equal(previousFilter, currentFilter) {
  300. changed = true
  301. }
  302. return nil
  303. })
  304. if err != nil {
  305. return false, errors.Trace(err)
  306. }
  307. return changed, nil
  308. }
  309. // ServerEntryIterator is used to iterate over
  310. // stored server entries in rank order.
  311. type ServerEntryIterator struct {
  312. config *Config
  313. applyServerAffinity bool
  314. serverEntryIDs [][]byte
  315. serverEntryIndex int
  316. isTacticsServerEntryIterator bool
  317. isTargetServerEntryIterator bool
  318. hasNextTargetServerEntry bool
  319. targetServerEntry *protocol.ServerEntry
  320. }
  321. // NewServerEntryIterator creates a new ServerEntryIterator.
  322. //
  323. // The boolean return value indicates whether to treat the first server(s)
  324. // as affinity servers or not. When the server entry selection filter changes
  325. // such as from a specific region to any region, or when there was no previous
  326. // filter/iterator, the the first server(s) are arbitrary and should not be
  327. // given affinity treatment.
  328. //
  329. // NewServerEntryIterator and any returned ServerEntryIterator are not
  330. // designed for concurrent use as not all related datastore operations are
  331. // performed in a single transaction.
  332. //
  333. func NewServerEntryIterator(config *Config) (bool, *ServerEntryIterator, error) {
  334. // When configured, this target server entry is the only candidate
  335. if config.TargetServerEntry != "" {
  336. return newTargetServerEntryIterator(config, false)
  337. }
  338. filterChanged, err := hasServerEntryFilterChanged(config)
  339. if err != nil {
  340. return false, nil, errors.Trace(err)
  341. }
  342. applyServerAffinity := !filterChanged
  343. iterator := &ServerEntryIterator{
  344. config: config,
  345. applyServerAffinity: applyServerAffinity,
  346. }
  347. err = iterator.reset(true)
  348. if err != nil {
  349. return false, nil, errors.Trace(err)
  350. }
  351. return applyServerAffinity, iterator, nil
  352. }
  353. func NewTacticsServerEntryIterator(config *Config) (*ServerEntryIterator, error) {
  354. // When configured, this target server entry is the only candidate
  355. if config.TargetServerEntry != "" {
  356. _, iterator, err := newTargetServerEntryIterator(config, true)
  357. return iterator, err
  358. }
  359. iterator := &ServerEntryIterator{
  360. config: config,
  361. isTacticsServerEntryIterator: true,
  362. }
  363. err := iterator.reset(true)
  364. if err != nil {
  365. return nil, errors.Trace(err)
  366. }
  367. return iterator, nil
  368. }
  369. // newTargetServerEntryIterator is a helper for initializing the TargetServerEntry case
  370. func newTargetServerEntryIterator(config *Config, isTactics bool) (bool, *ServerEntryIterator, error) {
  371. serverEntry, err := protocol.DecodeServerEntry(
  372. config.TargetServerEntry, config.loadTimestamp, protocol.SERVER_ENTRY_SOURCE_TARGET)
  373. if err != nil {
  374. return false, nil, errors.Trace(err)
  375. }
  376. if serverEntry.Tag == "" {
  377. serverEntry.Tag = protocol.GenerateServerEntryTag(
  378. serverEntry.IpAddress, serverEntry.WebServerSecret)
  379. }
  380. if isTactics {
  381. if len(serverEntry.GetSupportedTacticsProtocols()) == 0 {
  382. return false, nil, errors.TraceNew("TargetServerEntry does not support tactics protocols")
  383. }
  384. } else {
  385. if config.EgressRegion != "" && serverEntry.Region != config.EgressRegion {
  386. return false, nil, errors.TraceNew("TargetServerEntry does not support EgressRegion")
  387. }
  388. limitTunnelProtocols := config.GetClientParameters().Get().TunnelProtocols(parameters.LimitTunnelProtocols)
  389. if len(limitTunnelProtocols) > 0 {
  390. // At the ServerEntryIterator level, only limitTunnelProtocols is applied;
  391. // excludeIntensive is handled higher up.
  392. if len(serverEntry.GetSupportedProtocols(
  393. conditionallyEnabledComponents{},
  394. config.UseUpstreamProxy(),
  395. limitTunnelProtocols,
  396. false)) == 0 {
  397. return false, nil, errors.TraceNew("TargetServerEntry does not support LimitTunnelProtocols")
  398. }
  399. }
  400. }
  401. iterator := &ServerEntryIterator{
  402. isTacticsServerEntryIterator: isTactics,
  403. isTargetServerEntryIterator: true,
  404. hasNextTargetServerEntry: true,
  405. targetServerEntry: serverEntry,
  406. }
  407. NoticeInfo("using TargetServerEntry: %s", serverEntry.GetDiagnosticID())
  408. return false, iterator, nil
  409. }
  410. // Reset a NewServerEntryIterator to the start of its cycle. The next
  411. // call to Next will return the first server entry.
  412. func (iterator *ServerEntryIterator) Reset() error {
  413. return iterator.reset(false)
  414. }
  415. func (iterator *ServerEntryIterator) reset(isInitialRound bool) error {
  416. iterator.Close()
  417. if iterator.isTargetServerEntryIterator {
  418. iterator.hasNextTargetServerEntry = true
  419. return nil
  420. }
  421. // BoltDB implementation note:
  422. // We don't keep a transaction open for the duration of the iterator
  423. // because this would expose the following semantics to consumer code:
  424. //
  425. // Read-only transactions and read-write transactions ... generally
  426. // shouldn't be opened simultaneously in the same goroutine. This can
  427. // cause a deadlock as the read-write transaction needs to periodically
  428. // re-map the data file but it cannot do so while a read-only
  429. // transaction is open.
  430. // (https://github.com/boltdb/bolt)
  431. //
  432. // So the underlying serverEntriesBucket could change after the serverEntryIDs
  433. // list is built.
  434. var serverEntryIDs [][]byte
  435. err := datastoreView(func(tx *datastoreTx) error {
  436. bucket := tx.bucket(datastoreKeyValueBucket)
  437. serverEntryIDs = make([][]byte, 0)
  438. shuffleHead := 0
  439. var affinityServerEntryID []byte
  440. // In the first round only, move any server affinity candiate to the
  441. // very first position.
  442. if isInitialRound &&
  443. iterator.applyServerAffinity {
  444. affinityServerEntryID = bucket.get(datastoreAffinityServerEntryIDKey)
  445. if affinityServerEntryID != nil {
  446. serverEntryIDs = append(serverEntryIDs, append([]byte(nil), affinityServerEntryID...))
  447. shuffleHead = 1
  448. }
  449. }
  450. bucket = tx.bucket(datastoreServerEntriesBucket)
  451. cursor := bucket.cursor()
  452. for key := cursor.firstKey(); key != nil; key = cursor.nextKey() {
  453. if affinityServerEntryID != nil {
  454. if bytes.Equal(affinityServerEntryID, key) {
  455. continue
  456. }
  457. }
  458. serverEntryIDs = append(serverEntryIDs, append([]byte(nil), key...))
  459. }
  460. cursor.close()
  461. // Randomly shuffle the entire list of server IDs, excluding the
  462. // server affinity candidate.
  463. for i := len(serverEntryIDs) - 1; i > shuffleHead-1; i-- {
  464. j := prng.Intn(i+1-shuffleHead) + shuffleHead
  465. serverEntryIDs[i], serverEntryIDs[j] = serverEntryIDs[j], serverEntryIDs[i]
  466. }
  467. // In the first round, or with some probability, move _potential_ replay
  468. // candidates to the front of the list (excepting the server affinity slot,
  469. // if any). This move is post-shuffle so the order is still randomized. To
  470. // save the memory overhead of unmarshalling all dial parameters, this
  471. // operation just moves any server with a dial parameter record to the
  472. // front. Whether the dial parameter remains valid for replay -- TTL,
  473. // tactics/config unchanged, etc. --- is checked later.
  474. //
  475. // TODO: move only up to parameters.ReplayCandidateCount to front?
  476. p := iterator.config.GetClientParameters().Get()
  477. if (isInitialRound || p.WeightedCoinFlip(parameters.ReplayLaterRoundMoveToFrontProbability)) &&
  478. p.Int(parameters.ReplayCandidateCount) != 0 {
  479. networkID := []byte(iterator.config.GetNetworkID())
  480. dialParamsBucket := tx.bucket(datastoreDialParametersBucket)
  481. i := shuffleHead
  482. j := len(serverEntryIDs) - 1
  483. for {
  484. for ; i < j; i++ {
  485. key := makeDialParametersKey(serverEntryIDs[i], networkID)
  486. if dialParamsBucket.get(key) == nil {
  487. break
  488. }
  489. }
  490. for ; i < j; j-- {
  491. key := makeDialParametersKey(serverEntryIDs[j], networkID)
  492. if dialParamsBucket.get(key) != nil {
  493. break
  494. }
  495. }
  496. if i < j {
  497. serverEntryIDs[i], serverEntryIDs[j] = serverEntryIDs[j], serverEntryIDs[i]
  498. i++
  499. j--
  500. } else {
  501. break
  502. }
  503. }
  504. }
  505. return nil
  506. })
  507. if err != nil {
  508. return errors.Trace(err)
  509. }
  510. iterator.serverEntryIDs = serverEntryIDs
  511. iterator.serverEntryIndex = 0
  512. return nil
  513. }
  514. // Close cleans up resources associated with a ServerEntryIterator.
  515. func (iterator *ServerEntryIterator) Close() {
  516. iterator.serverEntryIDs = nil
  517. iterator.serverEntryIndex = 0
  518. }
  519. // Next returns the next server entry, by rank, for a ServerEntryIterator.
  520. // Returns nil with no error when there is no next item.
  521. func (iterator *ServerEntryIterator) Next() (*protocol.ServerEntry, error) {
  522. var serverEntry *protocol.ServerEntry
  523. var err error
  524. defer func() {
  525. if err != nil {
  526. iterator.Close()
  527. }
  528. }()
  529. if iterator.isTargetServerEntryIterator {
  530. if iterator.hasNextTargetServerEntry {
  531. iterator.hasNextTargetServerEntry = false
  532. return MakeCompatibleServerEntry(iterator.targetServerEntry), nil
  533. }
  534. return nil, nil
  535. }
  536. // There are no region/protocol indexes for the server entries bucket.
  537. // Loop until we have the next server entry that matches the iterator
  538. // filter requirements.
  539. for {
  540. if iterator.serverEntryIndex >= len(iterator.serverEntryIDs) {
  541. // There is no next item
  542. return nil, nil
  543. }
  544. serverEntryID := iterator.serverEntryIDs[iterator.serverEntryIndex]
  545. iterator.serverEntryIndex += 1
  546. serverEntry = nil
  547. err = datastoreView(func(tx *datastoreTx) error {
  548. serverEntries := tx.bucket(datastoreServerEntriesBucket)
  549. value := serverEntries.get(serverEntryID)
  550. if value == nil {
  551. return nil
  552. }
  553. // Must unmarshal here as slice is only valid within transaction.
  554. err = json.Unmarshal(value, &serverEntry)
  555. if err != nil {
  556. // In case of data corruption or a bug causing this condition,
  557. // do not stop iterating.
  558. serverEntry = nil
  559. NoticeWarning(
  560. "ServerEntryIterator.Next: json.Unmarshal failed: %s",
  561. errors.Trace(err))
  562. }
  563. return nil
  564. })
  565. if err != nil {
  566. return nil, errors.Trace(err)
  567. }
  568. if serverEntry == nil {
  569. // In case of data corruption or a bug causing this condition,
  570. // do not stop iterating.
  571. NoticeWarning("ServerEntryIterator.Next: unexpected missing server entry")
  572. continue
  573. }
  574. // Generate a derived server entry tag for server entries with no tag. Store
  575. // back the updated server entry so that (a) the tag doesn't need to be
  576. // regenerated; (b) the server entry can be looked up by tag (currently used
  577. // in the status request prune case).
  578. //
  579. // This is a distinct transaction so as to avoid the overhead of regular
  580. // write transactions in the iterator; once tags have been stored back, most
  581. // iterator transactions will remain read-only.
  582. if serverEntry.Tag == "" {
  583. serverEntry.Tag = protocol.GenerateServerEntryTag(
  584. serverEntry.IpAddress, serverEntry.WebServerSecret)
  585. err = datastoreUpdate(func(tx *datastoreTx) error {
  586. serverEntries := tx.bucket(datastoreServerEntriesBucket)
  587. serverEntryTags := tx.bucket(datastoreServerEntryTagsBucket)
  588. // We must reload and store back the server entry _fields_ to preserve any
  589. // currently unrecognized fields, for future compatibility.
  590. value := serverEntries.get(serverEntryID)
  591. if value == nil {
  592. return nil
  593. }
  594. var serverEntryFields protocol.ServerEntryFields
  595. err := json.Unmarshal(value, &serverEntryFields)
  596. if err != nil {
  597. return errors.Trace(err)
  598. }
  599. // As there is minor race condition between loading/checking serverEntry
  600. // and reloading/modifying serverEntryFields, this transaction references
  601. // only the freshly loaded fields when checking and setting the tag.
  602. serverEntryTag := serverEntryFields.GetTag()
  603. if serverEntryTag != "" {
  604. return nil
  605. }
  606. serverEntryTag = protocol.GenerateServerEntryTag(
  607. serverEntryFields.GetIPAddress(),
  608. serverEntryFields.GetWebServerSecret())
  609. serverEntryFields.SetTag(serverEntryTag)
  610. jsonServerEntryFields, err := json.Marshal(serverEntryFields)
  611. if err != nil {
  612. return errors.Trace(err)
  613. }
  614. serverEntries.put(serverEntryID, jsonServerEntryFields)
  615. if err != nil {
  616. return errors.Trace(err)
  617. }
  618. serverEntryTags.put([]byte(serverEntryTag), serverEntryID)
  619. if err != nil {
  620. return errors.Trace(err)
  621. }
  622. return nil
  623. })
  624. if err != nil {
  625. // Do not stop.
  626. NoticeWarning(
  627. "ServerEntryIterator.Next: update server entry failed: %s",
  628. errors.Trace(err))
  629. }
  630. }
  631. if iterator.serverEntryIndex%datastoreServerEntryFetchGCThreshold == 0 {
  632. DoGarbageCollection()
  633. }
  634. // Check filter requirements
  635. if iterator.isTacticsServerEntryIterator {
  636. // Tactics doesn't filter by egress region.
  637. if len(serverEntry.GetSupportedTacticsProtocols()) > 0 {
  638. break
  639. }
  640. } else {
  641. if iterator.config.EgressRegion == "" ||
  642. serverEntry.Region == iterator.config.EgressRegion {
  643. break
  644. }
  645. }
  646. }
  647. return MakeCompatibleServerEntry(serverEntry), nil
  648. }
  649. // MakeCompatibleServerEntry provides backwards compatibility with old server entries
  650. // which have a single meekFrontingDomain and not a meekFrontingAddresses array.
  651. // By copying this one meekFrontingDomain into meekFrontingAddresses, this client effectively
  652. // uses that single value as legacy clients do.
  653. func MakeCompatibleServerEntry(serverEntry *protocol.ServerEntry) *protocol.ServerEntry {
  654. if len(serverEntry.MeekFrontingAddresses) == 0 && serverEntry.MeekFrontingDomain != "" {
  655. serverEntry.MeekFrontingAddresses =
  656. append(serverEntry.MeekFrontingAddresses, serverEntry.MeekFrontingDomain)
  657. }
  658. return serverEntry
  659. }
  660. // PruneServerEntry deletes the server entry, along with associated data,
  661. // corresponding to the specified server entry tag. Pruning is subject to an
  662. // age check. In the case of an error, a notice is emitted.
  663. func PruneServerEntry(config *Config, serverEntryTag string) {
  664. err := pruneServerEntry(config, serverEntryTag)
  665. if err != nil {
  666. NoticeWarning(
  667. "PruneServerEntry failed: %s: %s",
  668. serverEntryTag, errors.Trace(err))
  669. return
  670. }
  671. NoticePruneServerEntry(serverEntryTag)
  672. }
  673. func pruneServerEntry(config *Config, serverEntryTag string) error {
  674. minimumAgeForPruning := config.GetClientParameters().Get().Duration(
  675. parameters.ServerEntryMinimumAgeForPruning)
  676. return datastoreUpdate(func(tx *datastoreTx) error {
  677. serverEntries := tx.bucket(datastoreServerEntriesBucket)
  678. serverEntryTags := tx.bucket(datastoreServerEntryTagsBucket)
  679. serverEntryTombstoneTags := tx.bucket(datastoreServerEntryTombstoneTagsBucket)
  680. keyValues := tx.bucket(datastoreKeyValueBucket)
  681. dialParameters := tx.bucket(datastoreDialParametersBucket)
  682. serverEntryTagBytes := []byte(serverEntryTag)
  683. serverEntryID := serverEntryTags.get(serverEntryTagBytes)
  684. if serverEntryID == nil {
  685. return errors.TraceNew("server entry tag not found")
  686. }
  687. serverEntryJson := serverEntries.get(serverEntryID)
  688. if serverEntryJson == nil {
  689. return errors.TraceNew("server entry not found")
  690. }
  691. var serverEntry *protocol.ServerEntry
  692. err := json.Unmarshal(serverEntryJson, &serverEntry)
  693. if err != nil {
  694. errors.Trace(err)
  695. }
  696. // Only prune sufficiently old server entries. This mitigates the case where
  697. // stale data in psiphond will incorrectly identify brand new servers as
  698. // being invalid/deleted.
  699. serverEntryLocalTimestamp, err := time.Parse(time.RFC3339, serverEntry.LocalTimestamp)
  700. if err != nil {
  701. errors.Trace(err)
  702. }
  703. if serverEntryLocalTimestamp.Add(minimumAgeForPruning).After(time.Now()) {
  704. return nil
  705. }
  706. // Handle the server IP recycle case where multiple serverEntryTags records
  707. // refer to the same server IP. Only delete the server entry record when its
  708. // tag matches the pruned tag. Otherwise, the server entry record is
  709. // associated with another tag. The pruned tag is still deleted.
  710. deleteServerEntry := (serverEntry.Tag == serverEntryTag)
  711. err = serverEntryTags.delete(serverEntryTagBytes)
  712. if err != nil {
  713. errors.Trace(err)
  714. }
  715. if deleteServerEntry {
  716. err = serverEntries.delete(serverEntryID)
  717. if err != nil {
  718. errors.Trace(err)
  719. }
  720. affinityServerEntryID := keyValues.get(datastoreAffinityServerEntryIDKey)
  721. if bytes.Equal(affinityServerEntryID, serverEntryID) {
  722. err = keyValues.delete(datastoreAffinityServerEntryIDKey)
  723. if err != nil {
  724. return errors.Trace(err)
  725. }
  726. }
  727. // TODO: expose boltdb Seek functionality to skip to first matching record.
  728. cursor := dialParameters.cursor()
  729. defer cursor.close()
  730. foundFirstMatch := false
  731. for key, _ := cursor.first(); key != nil; key, _ = cursor.next() {
  732. // Dial parameters key has serverID as a prefix; see makeDialParametersKey.
  733. if bytes.HasPrefix(key, serverEntryID) {
  734. foundFirstMatch = true
  735. err := dialParameters.delete(key)
  736. if err != nil {
  737. return errors.Trace(err)
  738. }
  739. } else if foundFirstMatch {
  740. break
  741. }
  742. }
  743. }
  744. // Tombstones prevent reimporting pruned server entries. Tombstone
  745. // identifiers are tags, which are derived from the web server secret in
  746. // addition to the server IP, so tombstones will not clobber recycled server
  747. // IPs as long as new web server secrets are generated in the recycle case.
  748. //
  749. // Tombstones are set only for embedded server entries, as all other sources
  750. // are expected to provide valid server entries; this also provides a fail-
  751. // safe mechanism to restore pruned server entries through all non-embedded
  752. // sources.
  753. if serverEntry.LocalSource == protocol.SERVER_ENTRY_SOURCE_EMBEDDED {
  754. err = serverEntryTombstoneTags.put(serverEntryTagBytes, []byte{1})
  755. if err != nil {
  756. return errors.Trace(err)
  757. }
  758. }
  759. return nil
  760. })
  761. }
  762. func scanServerEntries(scanner func(*protocol.ServerEntry)) error {
  763. err := datastoreView(func(tx *datastoreTx) error {
  764. bucket := tx.bucket(datastoreServerEntriesBucket)
  765. cursor := bucket.cursor()
  766. n := 0
  767. for key, value := cursor.first(); key != nil; key, value = cursor.next() {
  768. var serverEntry *protocol.ServerEntry
  769. err := json.Unmarshal(value, &serverEntry)
  770. if err != nil {
  771. // In case of data corruption or a bug causing this condition,
  772. // do not stop iterating.
  773. NoticeWarning("scanServerEntries: %s", errors.Trace(err))
  774. continue
  775. }
  776. scanner(serverEntry)
  777. n += 1
  778. if n == datastoreServerEntryFetchGCThreshold {
  779. DoGarbageCollection()
  780. n = 0
  781. }
  782. }
  783. cursor.close()
  784. return nil
  785. })
  786. if err != nil {
  787. return errors.Trace(err)
  788. }
  789. return nil
  790. }
  791. // CountServerEntries returns a count of stored server entries.
  792. func CountServerEntries() int {
  793. count := 0
  794. err := scanServerEntries(func(_ *protocol.ServerEntry) {
  795. count += 1
  796. })
  797. if err != nil {
  798. NoticeWarning("CountServerEntries failed: %s", err)
  799. return 0
  800. }
  801. return count
  802. }
  803. // CountServerEntriesWithConstraints returns a count of stored server entries for
  804. // the specified region and tunnel protocol limits.
  805. func CountServerEntriesWithConstraints(
  806. useUpstreamProxy bool,
  807. region string,
  808. constraints *protocolSelectionConstraints) (int, int) {
  809. // When CountServerEntriesWithConstraints is called only
  810. // limitTunnelProtocolState is fixed; excludeIntensive is transitory.
  811. excludeIntensive := false
  812. initialCount := 0
  813. count := 0
  814. err := scanServerEntries(func(serverEntry *protocol.ServerEntry) {
  815. if region == "" || serverEntry.Region == region {
  816. if constraints.isInitialCandidate(excludeIntensive, serverEntry) {
  817. initialCount += 1
  818. }
  819. if constraints.isCandidate(excludeIntensive, serverEntry) {
  820. count += 1
  821. }
  822. }
  823. })
  824. if err != nil {
  825. NoticeWarning("CountServerEntriesWithConstraints failed: %s", err)
  826. return 0, 0
  827. }
  828. return initialCount, count
  829. }
  830. // ReportAvailableRegions prints a notice with the available egress regions.
  831. // When limitState has initial protocols, the available regions are limited
  832. // to those available for the initial protocols; or if limitState has general
  833. // limited protocols, the available regions are similarly limited.
  834. func ReportAvailableRegions(config *Config, constraints *protocolSelectionConstraints) {
  835. // When ReportAvailableRegions is called only limitTunnelProtocolState is
  836. // fixed; excludeIntensive is transitory.
  837. excludeIntensive := false
  838. regions := make(map[string]bool)
  839. err := scanServerEntries(func(serverEntry *protocol.ServerEntry) {
  840. isCandidate := false
  841. if constraints.hasInitialProtocols() {
  842. isCandidate = constraints.isInitialCandidate(excludeIntensive, serverEntry)
  843. } else {
  844. isCandidate = constraints.isCandidate(excludeIntensive, serverEntry)
  845. }
  846. if isCandidate {
  847. regions[serverEntry.Region] = true
  848. }
  849. })
  850. if err != nil {
  851. NoticeWarning("ReportAvailableRegions failed: %s", err)
  852. return
  853. }
  854. regionList := make([]string, 0, len(regions))
  855. for region := range regions {
  856. // Some server entries do not have a region, but it makes no sense to return
  857. // an empty string as an "available region".
  858. if region != "" {
  859. regionList = append(regionList, region)
  860. }
  861. }
  862. NoticeAvailableEgressRegions(regionList)
  863. }
  864. // SetSplitTunnelRoutes updates the cached routes data for
  865. // the given region. The associated etag is also stored and
  866. // used to make efficient web requests for updates to the data.
  867. func SetSplitTunnelRoutes(region, etag string, data []byte) error {
  868. err := datastoreUpdate(func(tx *datastoreTx) error {
  869. bucket := tx.bucket(datastoreSplitTunnelRouteETagsBucket)
  870. err := bucket.put([]byte(region), []byte(etag))
  871. if err != nil {
  872. return errors.Trace(err)
  873. }
  874. bucket = tx.bucket(datastoreSplitTunnelRouteDataBucket)
  875. err = bucket.put([]byte(region), data)
  876. if err != nil {
  877. return errors.Trace(err)
  878. }
  879. return nil
  880. })
  881. if err != nil {
  882. return errors.Trace(err)
  883. }
  884. return nil
  885. }
  886. // GetSplitTunnelRoutesETag retrieves the etag for cached routes
  887. // data for the specified region. If not found, it returns an empty string value.
  888. func GetSplitTunnelRoutesETag(region string) (string, error) {
  889. var etag string
  890. err := datastoreView(func(tx *datastoreTx) error {
  891. bucket := tx.bucket(datastoreSplitTunnelRouteETagsBucket)
  892. etag = string(bucket.get([]byte(region)))
  893. return nil
  894. })
  895. if err != nil {
  896. return "", errors.Trace(err)
  897. }
  898. return etag, nil
  899. }
  900. // GetSplitTunnelRoutesData retrieves the cached routes data
  901. // for the specified region. If not found, it returns a nil value.
  902. func GetSplitTunnelRoutesData(region string) ([]byte, error) {
  903. var data []byte
  904. err := datastoreView(func(tx *datastoreTx) error {
  905. bucket := tx.bucket(datastoreSplitTunnelRouteDataBucket)
  906. value := bucket.get([]byte(region))
  907. if value != nil {
  908. // Must make a copy as slice is only valid within transaction.
  909. data = make([]byte, len(value))
  910. copy(data, value)
  911. }
  912. return nil
  913. })
  914. if err != nil {
  915. return nil, errors.Trace(err)
  916. }
  917. return data, nil
  918. }
  919. // SetUrlETag stores an ETag for the specfied URL.
  920. // Note: input URL is treated as a string, and is not
  921. // encoded or decoded or otherwise canonicalized.
  922. func SetUrlETag(url, etag string) error {
  923. err := datastoreUpdate(func(tx *datastoreTx) error {
  924. bucket := tx.bucket(datastoreUrlETagsBucket)
  925. err := bucket.put([]byte(url), []byte(etag))
  926. if err != nil {
  927. return errors.Trace(err)
  928. }
  929. return nil
  930. })
  931. if err != nil {
  932. return errors.Trace(err)
  933. }
  934. return nil
  935. }
  936. // GetUrlETag retrieves a previously stored an ETag for the
  937. // specfied URL. If not found, it returns an empty string value.
  938. func GetUrlETag(url string) (string, error) {
  939. var etag string
  940. err := datastoreView(func(tx *datastoreTx) error {
  941. bucket := tx.bucket(datastoreUrlETagsBucket)
  942. etag = string(bucket.get([]byte(url)))
  943. return nil
  944. })
  945. if err != nil {
  946. return "", errors.Trace(err)
  947. }
  948. return etag, nil
  949. }
  950. // SetKeyValue stores a key/value pair.
  951. func SetKeyValue(key, value string) error {
  952. err := datastoreUpdate(func(tx *datastoreTx) error {
  953. bucket := tx.bucket(datastoreKeyValueBucket)
  954. err := bucket.put([]byte(key), []byte(value))
  955. if err != nil {
  956. return errors.Trace(err)
  957. }
  958. return nil
  959. })
  960. if err != nil {
  961. return errors.Trace(err)
  962. }
  963. return nil
  964. }
  965. // GetKeyValue retrieves the value for a given key. If not found,
  966. // it returns an empty string value.
  967. func GetKeyValue(key string) (string, error) {
  968. var value string
  969. err := datastoreView(func(tx *datastoreTx) error {
  970. bucket := tx.bucket(datastoreKeyValueBucket)
  971. value = string(bucket.get([]byte(key)))
  972. return nil
  973. })
  974. if err != nil {
  975. return "", errors.Trace(err)
  976. }
  977. return value, nil
  978. }
  979. // Persistent stat records in the persistentStatStateUnreported
  980. // state are available for take out.
  981. //
  982. // Records in the persistentStatStateReporting have been taken
  983. // out and are pending either deletion (for a successful request)
  984. // or change to StateUnreported (for a failed request).
  985. //
  986. // All persistent stat records are reverted to StateUnreported
  987. // when the datastore is initialized at start up.
  988. var persistentStatStateUnreported = []byte("0")
  989. var persistentStatStateReporting = []byte("1")
  990. var persistentStatTypes = []string{
  991. datastorePersistentStatTypeRemoteServerList,
  992. datastorePersistentStatTypeFailedTunnel,
  993. }
  994. // StorePersistentStat adds a new persistent stat record, which
  995. // is set to StateUnreported and is an immediate candidate for
  996. // reporting.
  997. //
  998. // The stat is a JSON byte array containing fields as
  999. // required by the Psiphon server API. It's assumed that the
  1000. // JSON value contains enough unique information for the value to
  1001. // function as a key in the key/value datastore.
  1002. //
  1003. // Only up to PersistentStatsMaxStoreRecords are stored. Once this
  1004. // limit is reached, new records are discarded.
  1005. func StorePersistentStat(config *Config, statType string, stat []byte) error {
  1006. if !common.Contains(persistentStatTypes, statType) {
  1007. return errors.Tracef("invalid persistent stat type: %s", statType)
  1008. }
  1009. maxStoreRecords := config.GetClientParameters().Get().Int(
  1010. parameters.PersistentStatsMaxStoreRecords)
  1011. err := datastoreUpdate(func(tx *datastoreTx) error {
  1012. bucket := tx.bucket([]byte(statType))
  1013. count := 0
  1014. cursor := bucket.cursor()
  1015. for key, _ := cursor.first(); key != nil; key, _ = cursor.next() {
  1016. count++
  1017. }
  1018. cursor.close()
  1019. // TODO: assuming newer metrics are more useful, replace oldest record
  1020. // instead of discarding?
  1021. if count >= maxStoreRecords {
  1022. // Silently discard.
  1023. return nil
  1024. }
  1025. err := bucket.put(stat, persistentStatStateUnreported)
  1026. if err != nil {
  1027. return errors.Trace(err)
  1028. }
  1029. return nil
  1030. })
  1031. if err != nil {
  1032. return errors.Trace(err)
  1033. }
  1034. return nil
  1035. }
  1036. // CountUnreportedPersistentStats returns the number of persistent
  1037. // stat records in StateUnreported.
  1038. func CountUnreportedPersistentStats() int {
  1039. unreported := 0
  1040. err := datastoreView(func(tx *datastoreTx) error {
  1041. for _, statType := range persistentStatTypes {
  1042. bucket := tx.bucket([]byte(statType))
  1043. cursor := bucket.cursor()
  1044. for key, value := cursor.first(); key != nil; key, value = cursor.next() {
  1045. if bytes.Equal(value, persistentStatStateUnreported) {
  1046. unreported++
  1047. }
  1048. }
  1049. cursor.close()
  1050. }
  1051. return nil
  1052. })
  1053. if err != nil {
  1054. NoticeWarning("CountUnreportedPersistentStats failed: %s", err)
  1055. return 0
  1056. }
  1057. return unreported
  1058. }
  1059. // TakeOutUnreportedPersistentStats returns persistent stats records that are
  1060. // in StateUnreported. At least one record, if present, will be returned and
  1061. // then additional records up to PersistentStatsMaxSendBytes. The records are
  1062. // set to StateReporting. If the records are successfully reported, clear them
  1063. // with ClearReportedPersistentStats. If the records are not successfully
  1064. // reported, restore them with PutBackUnreportedPersistentStats.
  1065. func TakeOutUnreportedPersistentStats(config *Config) (map[string][][]byte, error) {
  1066. stats := make(map[string][][]byte)
  1067. maxSendBytes := config.GetClientParameters().Get().Int(
  1068. parameters.PersistentStatsMaxSendBytes)
  1069. err := datastoreUpdate(func(tx *datastoreTx) error {
  1070. sendBytes := 0
  1071. for _, statType := range persistentStatTypes {
  1072. bucket := tx.bucket([]byte(statType))
  1073. cursor := bucket.cursor()
  1074. for key, value := cursor.first(); key != nil; key, value = cursor.next() {
  1075. // Perform a test JSON unmarshaling. In case of data corruption or a bug,
  1076. // delete and skip the record.
  1077. var jsonData interface{}
  1078. err := json.Unmarshal(key, &jsonData)
  1079. if err != nil {
  1080. NoticeWarning(
  1081. "Invalid key in TakeOutUnreportedPersistentStats: %s: %s",
  1082. string(key), err)
  1083. bucket.delete(key)
  1084. continue
  1085. }
  1086. if bytes.Equal(value, persistentStatStateUnreported) {
  1087. // Must make a copy as slice is only valid within transaction.
  1088. data := make([]byte, len(key))
  1089. copy(data, key)
  1090. if stats[statType] == nil {
  1091. stats[statType] = make([][]byte, 0)
  1092. }
  1093. stats[statType] = append(stats[statType], data)
  1094. sendBytes += len(data)
  1095. if sendBytes >= maxSendBytes {
  1096. break
  1097. }
  1098. }
  1099. }
  1100. cursor.close()
  1101. for _, key := range stats[statType] {
  1102. err := bucket.put(key, persistentStatStateReporting)
  1103. if err != nil {
  1104. return errors.Trace(err)
  1105. }
  1106. }
  1107. }
  1108. return nil
  1109. })
  1110. if err != nil {
  1111. return nil, errors.Trace(err)
  1112. }
  1113. return stats, nil
  1114. }
  1115. // PutBackUnreportedPersistentStats restores a list of persistent
  1116. // stat records to StateUnreported.
  1117. func PutBackUnreportedPersistentStats(stats map[string][][]byte) error {
  1118. err := datastoreUpdate(func(tx *datastoreTx) error {
  1119. for _, statType := range persistentStatTypes {
  1120. bucket := tx.bucket([]byte(statType))
  1121. for _, key := range stats[statType] {
  1122. err := bucket.put(key, persistentStatStateUnreported)
  1123. if err != nil {
  1124. return errors.Trace(err)
  1125. }
  1126. }
  1127. }
  1128. return nil
  1129. })
  1130. if err != nil {
  1131. return errors.Trace(err)
  1132. }
  1133. return nil
  1134. }
  1135. // ClearReportedPersistentStats deletes a list of persistent
  1136. // stat records that were successfully reported.
  1137. func ClearReportedPersistentStats(stats map[string][][]byte) error {
  1138. err := datastoreUpdate(func(tx *datastoreTx) error {
  1139. for _, statType := range persistentStatTypes {
  1140. bucket := tx.bucket([]byte(statType))
  1141. for _, key := range stats[statType] {
  1142. err := bucket.delete(key)
  1143. if err != nil {
  1144. return err
  1145. }
  1146. }
  1147. }
  1148. return nil
  1149. })
  1150. if err != nil {
  1151. return errors.Trace(err)
  1152. }
  1153. return nil
  1154. }
  1155. // resetAllPersistentStatsToUnreported sets all persistent stat
  1156. // records to StateUnreported. This reset is called when the
  1157. // datastore is initialized at start up, as we do not know if
  1158. // persistent records in StateReporting were reported or not.
  1159. func resetAllPersistentStatsToUnreported() error {
  1160. err := datastoreUpdate(func(tx *datastoreTx) error {
  1161. for _, statType := range persistentStatTypes {
  1162. bucket := tx.bucket([]byte(statType))
  1163. resetKeys := make([][]byte, 0)
  1164. cursor := bucket.cursor()
  1165. for key := cursor.firstKey(); key != nil; key = cursor.nextKey() {
  1166. resetKeys = append(resetKeys, key)
  1167. }
  1168. cursor.close()
  1169. // TODO: data mutation is done outside cursor. Is this
  1170. // strictly necessary in this case? As is, this means
  1171. // all stats need to be loaded into memory at once.
  1172. // https://godoc.org/github.com/boltdb/bolt#Cursor
  1173. for _, key := range resetKeys {
  1174. err := bucket.put(key, persistentStatStateUnreported)
  1175. if err != nil {
  1176. return errors.Trace(err)
  1177. }
  1178. }
  1179. }
  1180. return nil
  1181. })
  1182. if err != nil {
  1183. return errors.Trace(err)
  1184. }
  1185. return nil
  1186. }
  1187. // CountSLOKs returns the total number of SLOK records.
  1188. func CountSLOKs() int {
  1189. count := 0
  1190. err := datastoreView(func(tx *datastoreTx) error {
  1191. bucket := tx.bucket(datastoreSLOKsBucket)
  1192. cursor := bucket.cursor()
  1193. for key := cursor.firstKey(); key != nil; key = cursor.nextKey() {
  1194. count++
  1195. }
  1196. cursor.close()
  1197. return nil
  1198. })
  1199. if err != nil {
  1200. NoticeWarning("CountSLOKs failed: %s", err)
  1201. return 0
  1202. }
  1203. return count
  1204. }
  1205. // DeleteSLOKs deletes all SLOK records.
  1206. func DeleteSLOKs() error {
  1207. err := datastoreUpdate(func(tx *datastoreTx) error {
  1208. return tx.clearBucket(datastoreSLOKsBucket)
  1209. })
  1210. if err != nil {
  1211. return errors.Trace(err)
  1212. }
  1213. return nil
  1214. }
  1215. // SetSLOK stores a SLOK key, referenced by its ID. The bool
  1216. // return value indicates whether the SLOK was already stored.
  1217. func SetSLOK(id, key []byte) (bool, error) {
  1218. var duplicate bool
  1219. err := datastoreUpdate(func(tx *datastoreTx) error {
  1220. bucket := tx.bucket(datastoreSLOKsBucket)
  1221. duplicate = bucket.get(id) != nil
  1222. err := bucket.put([]byte(id), []byte(key))
  1223. if err != nil {
  1224. return errors.Trace(err)
  1225. }
  1226. return nil
  1227. })
  1228. if err != nil {
  1229. return false, errors.Trace(err)
  1230. }
  1231. return duplicate, nil
  1232. }
  1233. // GetSLOK returns a SLOK key for the specified ID. The return
  1234. // value is nil if the SLOK is not found.
  1235. func GetSLOK(id []byte) ([]byte, error) {
  1236. var key []byte
  1237. err := datastoreView(func(tx *datastoreTx) error {
  1238. bucket := tx.bucket(datastoreSLOKsBucket)
  1239. key = bucket.get(id)
  1240. return nil
  1241. })
  1242. if err != nil {
  1243. return nil, errors.Trace(err)
  1244. }
  1245. return key, nil
  1246. }
  1247. func makeDialParametersKey(serverIPAddress, networkID []byte) []byte {
  1248. // TODO: structured key?
  1249. return append(append([]byte(nil), serverIPAddress...), networkID...)
  1250. }
  1251. // SetDialParameters stores dial parameters associated with the specified
  1252. // server/network ID.
  1253. func SetDialParameters(serverIPAddress, networkID string, dialParams *DialParameters) error {
  1254. key := makeDialParametersKey([]byte(serverIPAddress), []byte(networkID))
  1255. data, err := json.Marshal(dialParams)
  1256. if err != nil {
  1257. return errors.Trace(err)
  1258. }
  1259. return setBucketValue(datastoreDialParametersBucket, key, data)
  1260. }
  1261. // GetDialParameters fetches any dial parameters associated with the specified
  1262. // server/network ID. Returns nil, nil when no record is found.
  1263. func GetDialParameters(serverIPAddress, networkID string) (*DialParameters, error) {
  1264. key := makeDialParametersKey([]byte(serverIPAddress), []byte(networkID))
  1265. data, err := getBucketValue(datastoreDialParametersBucket, key)
  1266. if err != nil {
  1267. return nil, errors.Trace(err)
  1268. }
  1269. if data == nil {
  1270. return nil, nil
  1271. }
  1272. var dialParams *DialParameters
  1273. err = json.Unmarshal(data, &dialParams)
  1274. if err != nil {
  1275. return nil, errors.Trace(err)
  1276. }
  1277. return dialParams, nil
  1278. }
  1279. // DeleteDialParameters clears any dial parameters associated with the
  1280. // specified server/network ID.
  1281. func DeleteDialParameters(serverIPAddress, networkID string) error {
  1282. key := makeDialParametersKey([]byte(serverIPAddress), []byte(networkID))
  1283. return deleteBucketValue(datastoreDialParametersBucket, key)
  1284. }
  1285. // TacticsStorer implements tactics.Storer.
  1286. type TacticsStorer struct {
  1287. }
  1288. func (t *TacticsStorer) SetTacticsRecord(networkID string, record []byte) error {
  1289. return setBucketValue(datastoreTacticsBucket, []byte(networkID), record)
  1290. }
  1291. func (t *TacticsStorer) GetTacticsRecord(networkID string) ([]byte, error) {
  1292. return getBucketValue(datastoreTacticsBucket, []byte(networkID))
  1293. }
  1294. func (t *TacticsStorer) SetSpeedTestSamplesRecord(networkID string, record []byte) error {
  1295. return setBucketValue(datastoreSpeedTestSamplesBucket, []byte(networkID), record)
  1296. }
  1297. func (t *TacticsStorer) GetSpeedTestSamplesRecord(networkID string) ([]byte, error) {
  1298. return getBucketValue(datastoreSpeedTestSamplesBucket, []byte(networkID))
  1299. }
  1300. // GetTacticsStorer creates a TacticsStorer.
  1301. func GetTacticsStorer() *TacticsStorer {
  1302. return &TacticsStorer{}
  1303. }
  1304. // GetAffinityServerEntryAndDialParameters fetches the current affinity server
  1305. // entry value and any corresponding dial parameters for the specified network
  1306. // ID. An error is returned when no affinity server is available. The
  1307. // DialParameter output may be nil when a server entry is found but has no
  1308. // dial parameters.
  1309. func GetAffinityServerEntryAndDialParameters(
  1310. networkID string) (protocol.ServerEntryFields, *DialParameters, error) {
  1311. var serverEntryFields protocol.ServerEntryFields
  1312. var dialParams *DialParameters
  1313. err := datastoreView(func(tx *datastoreTx) error {
  1314. keyValues := tx.bucket(datastoreKeyValueBucket)
  1315. serverEntries := tx.bucket(datastoreServerEntriesBucket)
  1316. dialParameters := tx.bucket(datastoreDialParametersBucket)
  1317. affinityServerEntryID := keyValues.get(datastoreAffinityServerEntryIDKey)
  1318. if affinityServerEntryID == nil {
  1319. return errors.TraceNew("no affinity server available")
  1320. }
  1321. serverEntryRecord := serverEntries.get(affinityServerEntryID)
  1322. if serverEntryRecord == nil {
  1323. return errors.TraceNew("affinity server entry not found")
  1324. }
  1325. err := json.Unmarshal(
  1326. serverEntryRecord,
  1327. &serverEntryFields)
  1328. if err != nil {
  1329. return errors.Trace(err)
  1330. }
  1331. dialParamsKey := makeDialParametersKey(
  1332. []byte(serverEntryFields.GetIPAddress()),
  1333. []byte(networkID))
  1334. dialParamsRecord := dialParameters.get(dialParamsKey)
  1335. if dialParamsRecord != nil {
  1336. err := json.Unmarshal(dialParamsRecord, &dialParams)
  1337. if err != nil {
  1338. return errors.Trace(err)
  1339. }
  1340. }
  1341. return nil
  1342. })
  1343. if err != nil {
  1344. return nil, nil, errors.Trace(err)
  1345. }
  1346. return serverEntryFields, dialParams, nil
  1347. }
  1348. func setBucketValue(bucket, key, value []byte) error {
  1349. err := datastoreUpdate(func(tx *datastoreTx) error {
  1350. bucket := tx.bucket(bucket)
  1351. err := bucket.put(key, value)
  1352. if err != nil {
  1353. return errors.Trace(err)
  1354. }
  1355. return nil
  1356. })
  1357. if err != nil {
  1358. return errors.Trace(err)
  1359. }
  1360. return nil
  1361. }
  1362. func getBucketValue(bucket, key []byte) ([]byte, error) {
  1363. var value []byte
  1364. err := datastoreView(func(tx *datastoreTx) error {
  1365. bucket := tx.bucket(bucket)
  1366. value = bucket.get(key)
  1367. return nil
  1368. })
  1369. if err != nil {
  1370. return nil, errors.Trace(err)
  1371. }
  1372. return value, nil
  1373. }
  1374. func deleteBucketValue(bucket, key []byte) error {
  1375. err := datastoreUpdate(func(tx *datastoreTx) error {
  1376. bucket := tx.bucket(bucket)
  1377. return bucket.delete(key)
  1378. })
  1379. if err != nil {
  1380. return errors.Trace(err)
  1381. }
  1382. return nil
  1383. }