We build too many walls and not enough bridges

internal/store/user_address.go

@@ -0,0 +1,216 @@
+// Copyright (c) 2020 Proton Technologies AG
+//
+// This file is part of ProtonMail Bridge.
+//
+// ProtonMail Bridge is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// ProtonMail Bridge is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with ProtonMail Bridge.  If not, see <https://www.gnu.org/licenses/>.
+
+package store
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
+	"github.com/pkg/errors"
+	bolt "go.etcd.io/bbolt"
+)
+
+// GetAddress returns the store address by given ID.
+func (store *Store) GetAddress(addressID string) (*Address, error) {
+	store.lock.RLock()
+	defer store.lock.RUnlock()
+
+	storeAddress, ok := store.addresses[addressID]
+	if !ok {
+		return nil, fmt.Errorf("addressID %v does not exist", addressID)
+	}
+
+	return storeAddress, nil
+}
+
+// RebuildMailboxes truncates all mailbox buckets and recreates them from the metadata bucket again.
+func (store *Store) RebuildMailboxes() (err error) {
+	store.lock.Lock()
+	defer store.lock.Unlock()
+
+	log.WithField("user", store.UserID()).Trace("Truncating mailboxes")
+
+	if err = store.truncateMailboxesBucket(); err != nil {
+		log.WithError(err).Error("Could not truncate mailboxes bucket")
+		return
+	}
+
+	if err = store.truncateAddressInfoBucket(); err != nil {
+		log.WithError(err).Error("Could not truncate address info bucket")
+		return
+	}
+
+	if err = store.init(false); err != nil {
+		log.WithError(err).Error("Could not init store")
+		return
+	}
+
+	if err := store.increaseMailboxesVersion(); err != nil {
+		log.WithError(err).Error("Could not increase structure version")
+		// Do not return here. The truncation was already done and mode
+		// was changed in DB so we need to sync so that users start to see
+		// messages and not block other operations.
+	}
+
+	log.WithField("user", store.UserID()).Trace("Rebuilding mailboxes")
+	store.triggerSync()
+	return nil
+}
+
+// createOrDeleteAddressesEvent creates address objects in the store for each necessary address
+// and deletes any address objects that shouldn't be there.
+// It doesn't do anything to addresses that are rightfully there.
+// It should only be called from the event loop.
+func (store *Store) createOrDeleteAddressesEvent() (err error) {
+	labels, err := store.initCounts()
+	if err != nil {
+		return errors.Wrap(err, "failed to initialise label counts")
+	}
+
+	addrInfo, err := store.GetAddressInfo()
+	if err != nil {
+		return errors.Wrap(err, "failed to get addresses and address IDs")
+	}
+
+	// We need at least one address to continue.
+	if len(addrInfo) < 1 {
+		return errors.New("no addresses to initialise")
+	}
+
+	// If in combined mode, we only need the user's primary address.
+	if store.addressMode == combinedMode {
+		addrInfo = addrInfo[:1]
+	}
+
+	// Go through all addresses that *should* be there.
+	for _, addr := range addrInfo {
+		if _, ok := store.addresses[addr.AddressID]; ok {
+			continue
+		}
+
+		// This address is missing so we create it.
+		if err = store.addAddress(addr.Address, addr.AddressID, labels); err != nil {
+			return errors.Wrap(err, "failed to add address to store")
+		}
+	}
+
+	// Go through all addresses that *should not* be there.
+	for _, addr := range store.addresses {
+		belongs := false
+
+		for _, a := range addrInfo {
+			if addr.addressID == a.AddressID {
+				belongs = true
+				break
+			}
+		}
+
+		if belongs {
+			continue
+		}
+
+		delete(store.addresses, addr.addressID)
+	}
+
+	return err
+}
+
+// truncateAddressInfoBucket removes the address info bucket.
+func (store *Store) truncateAddressInfoBucket() (err error) {
+	log.Trace("Truncating address info bucket")
+
+	tx := func(tx *bolt.Tx) (err error) {
+		if err = tx.DeleteBucket(addressInfoBucket); err != nil {
+			return
+		}
+
+		if _, err = tx.CreateBucketIfNotExists(addressInfoBucket); err != nil {
+			return
+		}
+
+		return
+	}
+
+	return store.db.Update(tx)
+}
+
+// truncateMailboxesBucket removes the mailboxes bucket.
+func (store *Store) truncateMailboxesBucket() (err error) {
+	log.Trace("Truncating mailboxes bucket")
+
+	store.addresses = nil
+
+	tx := func(tx *bolt.Tx) (err error) {
+		mbs := tx.Bucket(mailboxesBucket)
+
+		return mbs.ForEach(func(addrIDMailbox, _ []byte) (err error) {
+			addr := mbs.Bucket(addrIDMailbox)
+
+			if err = addr.DeleteBucket(imapIDsBucket); err != nil {
+				return
+			}
+
+			if _, err = addr.CreateBucketIfNotExists(imapIDsBucket); err != nil {
+				return
+			}
+
+			if err = addr.DeleteBucket(apiIDsBucket); err != nil {
+				return
+			}
+
+			if _, err = addr.CreateBucketIfNotExists(apiIDsBucket); err != nil {
+				return
+			}
+
+			return
+		})
+	}
+
+	return store.db.Update(tx)
+}
+
+// initMailboxesBucket recreates the mailboxes bucket from the metadata bucket.
+func (store *Store) initMailboxesBucket() error { //nolint[unused]
+	i := 0
+	tx := func(tx *bolt.Tx) error {
+		return tx.Bucket(metadataBucket).ForEach(func(k, v []byte) error {
+			msg := &pmapi.Message{}
+
+			if err := json.Unmarshal(v, msg); err != nil {
+				return err
+			}
+
+			for _, a := range store.addresses {
+				if err := a.txCreateOrUpdateMessages(tx, []*pmapi.Message{msg}); err != nil {
+					return err
+				}
+			}
+
+			i++
+			if i%100 == 0 {
+				store.log.WithField("i", i).
+					Trace("Init mboxes heartbeat")
+			}
+
+			return nil
+		})
+	}
+
+	return store.db.Update(tx)
+}