Skip to content

Commit

Permalink
feat(GODT-1570): Badger store implementation
Browse files Browse the repository at this point in the history
Add implementation of Badger (https://github.com/dgraph-io/badger) as
an alternative cache implementation. This will become useful later as
part of GODT-1642 where we need to ensure that cache modifications can
be rolled back with failed db modifications and vice-versa.

The unit tests have been updated to use this new implementation and it
is also available to gluon-bench using the `-store=badger` option.

Finally the `StoreBuilder` has been updated to accept a slice of bytes
rather than a string as the encryption key.
  • Loading branch information
LBeernaertProton committed Aug 10, 2022
1 parent b57cfc9 commit 9caec09
Show file tree
Hide file tree
Showing 13 changed files with 363 additions and 9 deletions.
5 changes: 4 additions & 1 deletion benchmarks/gluon_bench/gluon_benchmarks/sync.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package gluon_benchmarks

import (
"context"
"crypto/sha256"
"flag"
"math/rand"
"time"
Expand Down Expand Up @@ -99,10 +100,12 @@ func (s *Sync) setupConnector(ctx context.Context) (utils.ConnectorImpl, error)
s.mailboxes = append(s.mailboxes, mboxID)
}

encryptionBytes := sha256.Sum256([]byte(*flags.UserPassword))

if _, err = s.server.AddUser(
ctx,
c.Connector(),
*flags.UserPassword); err != nil {
encryptionBytes[:]); err != nil {
return nil, err
}

Expand Down
5 changes: 4 additions & 1 deletion benchmarks/gluon_bench/imap_benchmarks/server/local.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package server

import (
"context"
"crypto/sha256"
"fmt"
"net"

Expand Down Expand Up @@ -72,10 +73,12 @@ func addUser(ctx context.Context, server *gluon.Server) error {
return nil
}

encryptionBytes := sha256.Sum256([]byte(*flags.UserPassword))

if userID, err := server.AddUser(
ctx,
c.Connector(),
*flags.UserPassword); err != nil {
encryptionBytes[:]); err != nil {
return err
} else if *flags.Verbose {
fmt.Printf("Adding user ID=%v\n", userID)
Expand Down
2 changes: 2 additions & 0 deletions benchmarks/gluon_bench/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,10 @@ import (
_ "github.com/ProtonMail/gluon/benchmarks/gluon_bench/gluon_benchmarks"
_ "github.com/ProtonMail/gluon/benchmarks/gluon_bench/imap_benchmarks"
_ "github.com/ProtonMail/gluon/benchmarks/gluon_bench/store_benchmarks"
"github.com/sirupsen/logrus"
)

func main() {
logrus.SetLevel(logrus.ErrorLevel)
benchmark.RunMain()
}
20 changes: 20 additions & 0 deletions benchmarks/gluon_bench/store_benchmarks/badger_store.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
package store_benchmarks

import (
"crypto/sha256"

"github.com/ProtonMail/gluon/benchmarks/gluon_bench/flags"
"github.com/ProtonMail/gluon/store"
"github.com/google/uuid"
)

type BadgerStoreBuilder struct{}

func (*BadgerStoreBuilder) New(path string) (store.Store, error) {
encryptionKey := sha256.Sum256([]byte(*flags.UserPassword))
return store.NewBadgerStore(path, uuid.NewString(), encryptionKey[:])
}

func init() {
RegisterStoreBuilder("badger", &BadgerStoreBuilder{})
}
2 changes: 1 addition & 1 deletion demo/demo.go
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ func addUser(ctx context.Context, server *gluon.Server, addresses []string, pass
userID, err := server.AddUser(
ctx,
connector,
password,
[]byte(password),
)
if err != nil {
return err
Expand Down
14 changes: 14 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ require (
entgo.io/ent v0.10.1
github.com/ProtonMail/gopenpgp/v2 v2.4.7
github.com/bradenaw/juniper v0.6.0
github.com/dgraph-io/badger/v3 v3.2103.2
github.com/emersion/go-imap v1.2.1-0.20220429085312-746087b7a317
github.com/emersion/go-imap-uidplus v0.0.0-20200503180755-e75854c361e9
github.com/emersion/go-mbox v1.0.2
Expand All @@ -27,19 +28,32 @@ require (
github.com/ProtonMail/go-mime v0.0.0-20220302105931-303f85f7fe0f // indirect
github.com/agext/levenshtein v1.2.1 // indirect
github.com/apparentlymart/go-textseg/v13 v13.0.0 // indirect
github.com/cespare/xxhash v1.1.0 // indirect
github.com/cespare/xxhash/v2 v2.1.1 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/dgraph-io/ristretto v0.1.0 // indirect
github.com/dustin/go-humanize v1.0.0 // indirect
github.com/emersion/go-sasl v0.0.0-20200509203442-7bfe0ed36a21 // indirect
github.com/go-openapi/inflect v0.19.0 // indirect
github.com/gogo/protobuf v1.3.2 // indirect
github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b // indirect
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
github.com/golang/protobuf v1.5.0 // indirect
github.com/golang/snappy v0.0.3 // indirect
github.com/google/flatbuffers v1.12.1 // indirect
github.com/google/go-cmp v0.5.6 // indirect
github.com/hashicorp/hcl/v2 v2.10.0 // indirect
github.com/klauspost/compress v1.12.3 // indirect
github.com/kr/pretty v0.3.0 // indirect
github.com/mitchellh/go-wordwrap v0.0.0-20150314170334-ad45545899c7 // indirect
github.com/pkg/errors v0.9.1 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/rogpeppe/go-internal v1.8.0 // indirect
github.com/zclconf/go-cty v1.8.0 // indirect
go.opencensus.io v0.23.0 // indirect
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519 // indirect
golang.org/x/mod v0.6.0-dev.0.20220106191415-9b9b3d81d5e3 // indirect
golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4 // indirect
golang.org/x/sys v0.0.0-20220712014510-0a85c31ab51e // indirect
gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
)
141 changes: 141 additions & 0 deletions go.sum

Large diffs are not rendered by default.

4 changes: 4 additions & 0 deletions internal/backend/user.go
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,10 @@ func (user *user) close(ctx context.Context) error {
// Wait until the connector update go routine has finished.
user.updateWG.Wait()

if err := user.store.Close(); err != nil {
return fmt.Errorf("failed to close user client storage: %w", err)
}

if err := user.client.Close(); err != nil {
return fmt.Errorf("failed to close user client: %w", err)
}
Expand Down
4 changes: 2 additions & 2 deletions server.go
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ func getDatabasePath(userPath, userID string) string {

// AddUser creates a new user and generates new unique ID for this user. If you have an existing userID, please use
// LoadUser instead.
func (s *Server) AddUser(ctx context.Context, conn connector.Connector, encryptionPassphrase string) (string, error) {
func (s *Server) AddUser(ctx context.Context, conn connector.Connector, encryptionPassphrase []byte) (string, error) {
userID := s.backend.NewUserID()

if err := s.LoadUser(ctx, conn, userID, encryptionPassphrase); err != nil {
Expand All @@ -108,7 +108,7 @@ func (s *Server) AddUser(ctx context.Context, conn connector.Connector, encrypti

// LoadUser loads an existing user's data from disk. This function can also be used to assign a custom userID to a mail
// server user.
func (s *Server) LoadUser(ctx context.Context, conn connector.Connector, userID, encryptionPassphrase string) error {
func (s *Server) LoadUser(ctx context.Context, conn connector.Connector, userID string, encryptionPassphrase []byte) error {
userPath, err := s.GetUserDataPath(userID)
if err != nil {
return err
Expand Down
157 changes: 157 additions & 0 deletions store/badger.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,157 @@
package store

import (
"path/filepath"
"sync"
"time"

"github.com/dgraph-io/badger/v3"
"github.com/sirupsen/logrus"
)

type BadgerStore struct {
db *badger.DB
gcExitCh chan struct{}
wg sync.WaitGroup
}

func logrusLevelToBadgerLevel(options badger.Options) badger.Options {
switch logrus.GetLevel() {
case logrus.InfoLevel:
return options.WithLoggingLevel(badger.INFO)
case logrus.TraceLevel:
fallthrough
case logrus.DebugLevel:
return options.WithLoggingLevel(badger.DEBUG)
case logrus.WarnLevel:
return options.WithLoggingLevel(badger.WARNING)
case logrus.FatalLevel:
fallthrough
case logrus.PanicLevel:
fallthrough
case logrus.ErrorLevel:
return options.WithLoggingLevel(badger.ERROR)
default:
return options.WithLoggingLevel(badger.ERROR)
}
}

func NewBadgerStore(path string, userID string, encryptionPassphrase []byte) (*BadgerStore, error) {
directory := filepath.Join(path, userID)
db, err := badger.Open(logrusLevelToBadgerLevel(badger.DefaultOptions(directory)).
WithLogger(logrus.StandardLogger()).
WithEncryptionKey(encryptionPassphrase).
WithIndexCacheSize(128 * 1024 * 1024))

if err != nil {
return nil, nil
}

store := &BadgerStore{db: db, gcExitCh: make(chan struct{})}

store.startGCCollector()

return store, nil
}

func (b *BadgerStore) startGCCollector() {
// Garbage collection needs to be run manually by us at some point.
// See https://dgraph.io/docs/badger/get-started/#garbage-collection for more details.
b.wg.Add(1)

go func() {
defer b.wg.Done()

gcRun := time.After(5 * time.Minute)

select {
case <-gcRun:
{
again:
err := b.db.RunValueLogGC(0.7)
if err == nil {
goto again
}
}
case <-b.gcExitCh:
return
}
}()
}

func (b *BadgerStore) Get(messageID string) ([]byte, error) {
var data []byte

if err := b.db.View(func(txn *badger.Txn) error {
item, err := txn.Get([]byte(messageID))
if err != nil {
return err
}

data, err = item.ValueCopy(nil)
if err != nil {
return err
}

return nil
}); err != nil {
return nil, err
}

return data, nil
}

func (b *BadgerStore) Set(messageID string, literal []byte) error {
return b.db.Update(func(txn *badger.Txn) error {
return txn.Set([]byte(messageID), literal)
})
}

func (b *BadgerStore) Update(oldID, newID string) error {
return b.db.Update(func(txn *badger.Txn) error {
oldIDBytes := []byte(oldID)
newIDBytes := []byte(newID)

item, err := txn.Get(oldIDBytes)
if err != nil {
return err
}

buffer := make([]byte, item.ValueSize())
buffer, err = item.ValueCopy(buffer)
if err != nil {
return err
}

if err := txn.Set(newIDBytes, buffer); err != nil {
return err
}

return txn.Delete(oldIDBytes)
})
}

func (b *BadgerStore) Delete(messageID ...string) error {
return b.db.Update(func(txn *badger.Txn) error {
for _, v := range messageID {
if err := txn.Delete([]byte(v)); err != nil {
return err
}
}

return nil
})
}

func (b *BadgerStore) Close() error {
close(b.gcExitCh)
b.wg.Wait()

return b.db.Close()
}

type BadgerStoreBuilder struct{}

func (*BadgerStoreBuilder) New(directory, userID, encryptionPassphrase string) (Store, error) {
return NewBadgerStore(directory, userID, []byte(encryptionPassphrase))
}
4 changes: 2 additions & 2 deletions store/disk.go
Original file line number Diff line number Diff line change
Expand Up @@ -131,8 +131,8 @@ func (c *onDiskStore) Close() error {

type OnDiskStoreBuilder struct{}

func (*OnDiskStoreBuilder) New(path, userID, userPassword string) (Store, error) {
func (*OnDiskStoreBuilder) New(path, userID string, userPassword []byte) (Store, error) {
storePath := filepath.Join(path, userID)

return NewOnDiskStore(storePath, []byte(userPassword))
return NewOnDiskStore(storePath, userPassword)
}
2 changes: 1 addition & 1 deletion store/store.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,5 +9,5 @@ type Store interface {
}

type StoreBuilder interface {
New(directory, userID, encryptionPassphrase string) (Store, error)
New(directory, userID string, encryptionPassphrase []byte) (Store, error)
}
12 changes: 11 additions & 1 deletion tests/server_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import (
"github.com/ProtonMail/gluon/connector"
"github.com/ProtonMail/gluon/imap"
"github.com/ProtonMail/gluon/internal"
"github.com/ProtonMail/gluon/store"
"github.com/emersion/go-imap/client"
"github.com/google/uuid"
"github.com/sirupsen/logrus"
Expand Down Expand Up @@ -177,6 +178,14 @@ func defaultServerOptions(tb testing.TB, modifiers ...serverOption) *serverOptio
return options
}

// Wrapper to ensure we always pass 32bytes worth of encryption key to the tests.
type testBadgerStoreBuilder struct{}

func (*testBadgerStoreBuilder) New(directory, userID string, encryptionPassphrase []byte) (store.Store, error) {
encryptionBytes := sha256.Sum256(encryptionPassphrase)
return store.NewBadgerStore(directory, userID, encryptionBytes[:])
}

// runServerWithPaths initializes and starts the mailserver using a pathGenerator.
func runServer(tb testing.TB, options *serverOptions, tests func(*testSession)) {
loggerIn := logrus.StandardLogger().WriterLevel(logrus.TraceLevel)
Expand All @@ -203,6 +212,7 @@ func runServer(tb testing.TB, options *serverOptions, tests func(*testSession))
gluon.WithVersionInfo(TestServerVersionInfo.Version.Major, TestServerVersionInfo.Version.Minor, TestServerVersionInfo.Version.Patch,
TestServerVersionInfo.Name, TestServerVersionInfo.Vendor, TestServerVersionInfo.SupportURL),
gluon.WithDataPath(storePath),
gluon.WithStoreBuilder(&testBadgerStoreBuilder{}),
)
require.NoError(tb, err)

Expand All @@ -227,7 +237,7 @@ func runServer(tb testing.TB, options *serverOptions, tests func(*testSession))
hash := sha256.Sum256([]byte(creds.usernames[0]))
userID := hex.EncodeToString(hash[:])

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

require.NoError(tb, conn.Sync(ctx))
Expand Down

0 comments on commit 9caec09

Please sign in to comment.