Skip to content

Commit 3a00849

Browse files
jameshoulahanjameshoulahan
authored
fix: Badger GC, compression, encryption (#135)
1 parent 38de1de commit 3a00849

7 files changed

Lines changed: 56 additions & 43 deletions

File tree

‎benchmarks/gluon_bench/gluon_benchmarks/sync.go‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ package gluon_benchmarks
22

33
import (
44
"context"
5-
"crypto/sha256"
65
"flag"
76
"math/rand"
87
"time"
@@ -14,6 +13,7 @@ import (
1413
"github.com/ProtonMail/gluon/benchmarks/gluon_bench/timing"
1514
"github.com/ProtonMail/gluon/benchmarks/gluon_bench/utils"
1615
"github.com/ProtonMail/gluon/imap"
16+
"github.com/ProtonMail/gluon/internal/hash"
1717
"github.com/google/uuid"
1818
_ "github.com/mattn/go-sqlite3"
1919
"github.com/sirupsen/logrus"
@@ -113,12 +113,11 @@ func (s *Sync) setupConnector(ctx context.Context) (utils.ConnectorImpl, error)
113113
s.mailboxes = append(s.mailboxes, mboxID)
114114
}
115115

116-
encryptionBytes := sha256.Sum256([]byte(*flags.UserPassword))
117-
118116
if _, err = s.server.AddUser(
119117
ctx,
120118
c.Connector(),
121-
encryptionBytes[:]); err != nil {
119+
hash.SHA256([]byte(*flags.UserPassword)),
120+
); err != nil {
122121
return nil, err
123122
}
124123

‎benchmarks/gluon_bench/imap_benchmarks/server/local.go‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,13 @@ package server
22

33
import (
44
"context"
5-
"crypto/sha256"
65
"fmt"
76
"net"
87

98
"github.com/ProtonMail/gluon"
109
"github.com/ProtonMail/gluon/benchmarks/gluon_bench/flags"
1110
"github.com/ProtonMail/gluon/benchmarks/gluon_bench/utils"
11+
"github.com/ProtonMail/gluon/internal/hash"
1212
"github.com/ProtonMail/gluon/profiling"
1313
_ "github.com/mattn/go-sqlite3"
1414
"github.com/sirupsen/logrus"
@@ -80,14 +80,16 @@ func addUser(ctx context.Context, server *gluon.Server) error {
8080
return err
8181
}
8282

83-
encryptionBytes := sha256.Sum256([]byte(*flags.UserPassword))
84-
85-
if userID, err := server.AddUser(
83+
userID, err := server.AddUser(
8684
ctx,
8785
c.Connector(),
88-
encryptionBytes[:]); err != nil {
86+
hash.SHA256([]byte(*flags.UserPassword)),
87+
)
88+
if err != nil {
8989
return err
90-
} else if *flags.Verbose {
90+
}
91+
92+
if *flags.Verbose {
9193
fmt.Printf("Adding user ID=%v\n", userID)
9294
}
9395

‎benchmarks/gluon_bench/store_benchmarks/badger_store.go‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,6 @@
11
package store_benchmarks
22

33
import (
4-
"crypto/sha256"
5-
64
"github.com/ProtonMail/gluon/benchmarks/gluon_bench/flags"
75
"github.com/ProtonMail/gluon/store"
86
"github.com/google/uuid"
@@ -11,8 +9,7 @@ import (
119
type BadgerStoreBuilder struct{}
1210

1311
func (*BadgerStoreBuilder) New(path string) (store.Store, error) {
14-
encryptionKey := sha256.Sum256([]byte(*flags.UserPassword))
15-
return store.NewBadgerStore(path, uuid.NewString(), encryptionKey[:])
12+
return store.NewBadgerStore(path, uuid.NewString(), []byte(*flags.UserPassword))
1613
}
1714

1815
func init() {

‎internal/hash/hash.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
package hash
2+
3+
import "crypto/sha256"
4+
5+
func SHA256(key []byte) []byte {
6+
hash := sha256.Sum256(key)
7+
8+
return hash[:]
9+
}

‎internal/utils/utils.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package utils
22

3-
import "github.com/google/uuid"
3+
import (
4+
"github.com/google/uuid"
5+
)
46

57
// NewRandomUserID return a new random user ID. For debugging purposes, the ID starts with the 'user-' prefix.
68
func NewRandomUserID() string {

‎store/badger.go‎

Lines changed: 30 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,41 +1,42 @@
11
package store
22

33
import (
4-
"crypto/sha256"
54
"os"
65
"path/filepath"
76
"sync"
87
"time"
98

109
"github.com/ProtonMail/gluon/imap"
10+
"github.com/ProtonMail/gluon/internal/hash"
1111
"github.com/dgraph-io/badger/v3"
1212
"github.com/sirupsen/logrus"
1313
)
1414

1515
type BadgerStore struct {
16-
db *badger.DB
17-
gcExitCh chan struct{}
18-
wg sync.WaitGroup
16+
db *badger.DB
17+
stopCh chan struct{}
18+
stopWG sync.WaitGroup
1919
}
2020

2121
type badgerTransaction struct {
2222
tx *badger.Txn
2323
}
2424

25-
func NewBadgerStore(path string, userID string, encryptionPassphrase []byte) (*BadgerStore, error) {
26-
encryptionKey := sha256.Sum256(encryptionPassphrase)
27-
25+
func NewBadgerStore(path string, userID string, passphrase []byte) (*BadgerStore, error) {
2826
db, err := badger.Open(badger.DefaultOptions(filepath.Join(path, userID)).
2927
WithLogger(logrus.StandardLogger()).
3028
WithLoggingLevel(badger.ERROR).
31-
WithEncryptionKey(encryptionKey[:]).
29+
WithEncryptionKey(hash.SHA256(passphrase)).
3230
WithIndexCacheSize(128 * 1024 * 1024),
3331
)
3432
if err != nil {
3533
return nil, err
3634
}
3735

38-
store := &BadgerStore{db: db, gcExitCh: make(chan struct{})}
36+
store := &BadgerStore{
37+
db: db,
38+
stopCh: make(chan struct{}),
39+
}
3940

4041
store.startGCCollector()
4142

@@ -45,24 +46,27 @@ func NewBadgerStore(path string, userID string, encryptionPassphrase []byte) (*B
4546
func (b *BadgerStore) startGCCollector() {
4647
// Garbage collection needs to be run manually by us at some point.
4748
// See https://dgraph.io/docs/badger/get-started/#garbage-collection for more details.
48-
b.wg.Add(1)
49+
b.stopWG.Add(1)
4950

5051
go func() {
51-
defer b.wg.Done()
52-
53-
gcRun := time.After(5 * time.Minute)
54-
55-
select {
56-
case <-gcRun:
57-
{
58-
again:
59-
err := b.db.RunValueLogGC(0.7)
60-
if err == nil {
61-
goto again
52+
defer b.stopWG.Done()
53+
54+
ticker := time.NewTicker(5 * time.Minute)
55+
defer ticker.Stop()
56+
57+
for {
58+
select {
59+
case <-ticker.C:
60+
{
61+
again:
62+
if err := b.db.RunValueLogGC(0.5); err == nil {
63+
goto again
64+
}
6265
}
66+
67+
case <-b.stopCh:
68+
return
6369
}
64-
case <-b.gcExitCh:
65-
return
6670
}
6771
}()
6872
}
@@ -118,8 +122,9 @@ func (b *badgerTransaction) Rollback() error {
118122
}
119123

120124
func (b *BadgerStore) Close() error {
121-
close(b.gcExitCh)
122-
b.wg.Wait()
125+
close(b.stopCh)
126+
127+
b.stopWG.Wait()
123128

124129
return b.db.Close()
125130
}

‎tests/server_test.go‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ package tests
22

33
import (
44
"context"
5-
"crypto/sha256"
65
"crypto/tls"
76
"encoding/hex"
87
"fmt"
@@ -16,6 +15,7 @@ import (
1615
"github.com/ProtonMail/gluon/connector"
1716
"github.com/ProtonMail/gluon/imap"
1817
"github.com/ProtonMail/gluon/internal"
18+
"github.com/ProtonMail/gluon/internal/hash"
1919
"github.com/ProtonMail/gluon/store"
2020
"github.com/emersion/go-imap/client"
2121
"github.com/google/uuid"
@@ -168,8 +168,7 @@ func runServer(tb testing.TB, options *serverOptions, tests func(session *testSe
168168
)
169169

170170
// Force USER ID to be consistent.
171-
hash := sha256.Sum256([]byte(creds.usernames[0]))
172-
userID := hex.EncodeToString(hash[:])
171+
userID := hex.EncodeToString(hash.SHA256([]byte(creds.usernames[0])))
173172

174173
err := server.LoadUser(ctx, conn, userID, []byte(creds.password))
175174
require.NoError(tb, err)

0 commit comments

Comments
 (0)