We build too many walls and not enough bridges

internal/store/user_address_info.go

@@ -0,0 +1,158 @@
+// 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"
+	"strings"
+
+	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
+	"github.com/pkg/errors"
+	bolt "go.etcd.io/bbolt"
+)
+
+// AddressInfo is the format of the data held in the addresses bucket in the store.
+// It allows us to easily keep an address and its ID together and serialisation/deserialisation to []byte.
+type AddressInfo struct {
+	Address, AddressID string
+}
+
+// GetAddressID returns the ID of the given address.
+func (store *Store) GetAddressID(addr string) (id string, err error) {
+	addrs, err := store.GetAddressInfo()
+	if err != nil {
+		return
+	}
+
+	for _, addrInfo := range addrs {
+		if strings.EqualFold(addrInfo.Address, addr) {
+			id = addrInfo.AddressID
+			return
+		}
+	}
+
+	err = errors.New("no such address")
+
+	return
+}
+
+// GetAddressInfo returns information about all addresses owned by the user.
+// The first element is the user's primary address and the rest (if present) are aliases.
+// It tries to source the information from the store but if the store doesn't yet have it, it
+// fetches it from the API and caches it for later.
+func (store *Store) GetAddressInfo() (addrs []AddressInfo, err error) {
+	if addrs, err = store.getAddressInfoFromStore(); err == nil && len(addrs) > 0 {
+		return
+	}
+
+	// Store does not have address info yet, need to build it first from API.
+	addressList := store.api.Addresses()
+	if addressList == nil {
+		err = errors.New("addresses unavailable")
+		store.log.WithError(err).Error("Could not get user addresses from API")
+		return
+	}
+
+	if err = store.createOrUpdateAddressInfo(addressList); err != nil {
+		store.log.WithError(err).Warn("Could not update address IDs in store")
+		return
+	}
+
+	return store.getAddressInfoFromStore()
+}
+
+// getAddressIDsByAddressFromStore returns a map from address to addressID for each address owned by the user.
+func (store *Store) getAddressInfoFromStore() (addrs []AddressInfo, err error) {
+	store.log.Debug("Retrieving address info from store")
+
+	tx := func(tx *bolt.Tx) (err error) {
+		c := tx.Bucket(addressInfoBucket).Cursor()
+		for index, addrInfoBytes := c.First(); index != nil; index, addrInfoBytes = c.Next() {
+			var addrInfo AddressInfo
+
+			if err = json.Unmarshal(addrInfoBytes, &addrInfo); err != nil {
+				store.log.WithError(err).Error("Could not unmarshal address and addressID")
+				return
+			}
+
+			addrs = append(addrs, addrInfo)
+		}
+
+		return
+	}
+
+	err = store.db.View(tx)
+
+	return
+}
+
+// createOrUpdateAddressInfo updates the store address/addressID bucket to match the given address list.
+// The address list supplied is assumed to contain active emails in any order.
+// It firstly (and stupidly) deletes the bucket of addresses and then fills it with up to date info.
+// This is because a user might delete an address and we don't want old addresses lying around (and finding the
+// specific ones to delete is likely not much more efficient than just rebuilding from scratch).
+func (store *Store) createOrUpdateAddressInfo(addressList pmapi.AddressList) (err error) {
+	tx := func(tx *bolt.Tx) error {
+		if err := tx.DeleteBucket(addressInfoBucket); err != nil {
+			store.log.WithError(err).Error("Could not delete addressIDs bucket")
+			return err
+		}
+
+		if _, err := tx.CreateBucketIfNotExists(addressInfoBucket); err != nil {
+			store.log.WithError(err).Error("Could not recreate addressIDs bucket")
+			return err
+		}
+
+		addrsBucket := tx.Bucket(addressInfoBucket)
+
+		for index, address := range filterAddresses(addressList) {
+			ib := itob(uint32(index))
+
+			info, err := json.Marshal(AddressInfo{
+				Address:   address.Email,
+				AddressID: address.ID,
+			})
+			if err != nil {
+				store.log.WithError(err).Error("Could not marshal address and addressID")
+				return err
+			}
+
+			if err := addrsBucket.Put(ib, info); err != nil {
+				store.log.WithError(err).Error("Could not put address and addressID into store")
+				return err
+			}
+		}
+
+		return nil
+	}
+
+	return store.db.Update(tx)
+}
+
+// filterAddresses filters out inactive addresses and ensures the original address is listed first.
+func filterAddresses(addressList pmapi.AddressList) (filteredList pmapi.AddressList) {
+	for _, address := range addressList {
+		if address.Receive != pmapi.CanReceive {
+			continue
+		}
+
+		filteredList = append(filteredList, address)
+	}
+
+	return
+}