From b9740e1b7d6fd99e4f83c8a79e636ea4a9c88cf3 Mon Sep 17 00:00:00 2001
From: Michal Horejsek <horejsek@protonmail.com>
Date: Tue, 3 Nov 2020 09:58:56 +0100
Subject: [PATCH] Close connection before deleting labels to prevent panics
 accessing deleted bucket

---
 Changelog.md                  |  1 +
 internal/store/mailbox.go     | 11 +++++++++++
 internal/store/mocks/mocks.go | 12 ++++++++++++
 internal/store/types.go       |  1 +
 internal/users/user.go        |  8 ++++----
 internal/users/users.go       |  2 +-
 pkg/pmapi/mocks/mocks.go      |  5 ++---
 7 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/Changelog.md b/Changelog.md
index 01c98d59..25137fae 100644
--- a/Changelog.md
+++ b/Changelog.md
@@ -15,6 +15,7 @@ Changelog [format](http://keepachangelog.com/en/1.0.0/)
 * GODT-849 Show in error counts in the end also lost messages.
 * GODT-835 Do not include conversation ID in references to show properly conversation threads in clients.
 * GODT-799 Fix skipping unwanted folders importing from mbox files.
+* GODT-769 Close connection before deleting labels to prevent panics accessing deleted bucket.
 
 ## [IE 1.2.0] Elbe
 
diff --git a/internal/store/mailbox.go b/internal/store/mailbox.go
index 89857af3..22d61aff 100644
--- a/internal/store/mailbox.go
+++ b/internal/store/mailbox.go
@@ -38,6 +38,8 @@ type Mailbox struct {
 	color       string
 
 	log *logrus.Entry
+
+	isDeleting bool
 }
 
 func newMailbox(storeAddress *Address, labelID, labelPrefix, labelName, color string) (mb *Mailbox, err error) {
@@ -215,6 +217,7 @@ func (storeMailbox *Mailbox) Rename(newName string) error {
 // Deletion has to be propagated to all the same mailboxes in all addresses.
 // The propagation is processed by the event loop.
 func (storeMailbox *Mailbox) Delete() error {
+	storeMailbox.isDeleting = true
 	return storeMailbox.storeAddress.deleteMailbox(storeMailbox.labelID)
 }
 
@@ -226,6 +229,14 @@ func (storeMailbox *Mailbox) GetDelimiter() string {
 // deleteMailboxEvent deletes the mailbox bucket.
 // This is called from the event loop.
 func (storeMailbox *Mailbox) deleteMailboxEvent() error {
+	if !storeMailbox.isDeleting {
+		// Deleting label removes bucket. Any ongoing connection selected
+		// in such mailbox then might panic because of non-existing bucket.
+		// Closing connetions prevents that panic but if the connection
+		// asked for deletion, it should not be closed so it can receive
+		// successful response.
+		storeMailbox.store.user.CloseAllConnections()
+	}
 	return storeMailbox.db().Update(func(tx *bolt.Tx) error {
 		return tx.Bucket(mailboxesBucket).DeleteBucket(storeMailbox.getBucketName())
 	})
diff --git a/internal/store/mocks/mocks.go b/internal/store/mocks/mocks.go
index 8fd5c2d6..ef208982 100644
--- a/internal/store/mocks/mocks.go
+++ b/internal/store/mocks/mocks.go
@@ -106,6 +106,18 @@ func (m *MockBridgeUser) EXPECT() *MockBridgeUserMockRecorder {
 	return m.recorder
 }
 
+// CloseAllConnections mocks base method
+func (m *MockBridgeUser) CloseAllConnections() {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "CloseAllConnections")
+}
+
+// CloseAllConnections indicates an expected call of CloseAllConnections
+func (mr *MockBridgeUserMockRecorder) CloseAllConnections() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CloseAllConnections", reflect.TypeOf((*MockBridgeUser)(nil).CloseAllConnections))
+}
+
 // CloseConnection mocks base method
 func (m *MockBridgeUser) CloseConnection(arg0 string) {
 	m.ctrl.T.Helper()
diff --git a/internal/store/types.go b/internal/store/types.go
index 6470c3dd..a23b458e 100644
--- a/internal/store/types.go
+++ b/internal/store/types.go
@@ -36,6 +36,7 @@ type BridgeUser interface {
 	GetPrimaryAddress() string
 	GetStoreAddresses() []string
 	UpdateUser() error
+	CloseAllConnections()
 	CloseConnection(string)
 	Logout() error
 }
diff --git a/internal/users/user.go b/internal/users/user.go
index 159988b5..b1e24e57 100644
--- a/internal/users/user.go
+++ b/internal/users/user.go
@@ -437,7 +437,7 @@ func (u *User) SwitchAddressMode() (err error) {
 
 	u.lock.Lock()
 	defer u.lock.Unlock()
-	u.closeAllConnections()
+	u.CloseAllConnections()
 
 	if u.store == nil {
 		err = errors.New("store is not initialised")
@@ -509,7 +509,7 @@ func (u *User) Logout() (err error) {
 	// Do not close whole store, just event loop. Some information might be needed offline (e.g. addressID)
 	u.closeEventLoop()
 
-	u.closeAllConnections()
+	u.CloseAllConnections()
 
 	runtime.GC()
 
@@ -532,8 +532,8 @@ func (u *User) closeEventLoop() {
 	u.store.CloseEventLoop()
 }
 
-// closeAllConnections calls CloseConnection for all users addresses.
-func (u *User) closeAllConnections() {
+// CloseAllConnections calls CloseConnection for all users addresses.
+func (u *User) CloseAllConnections() {
 	for _, address := range u.creds.EmailList() {
 		u.CloseConnection(address)
 	}
diff --git a/internal/users/users.go b/internal/users/users.go
index 62c6cf20..7cf4a318 100644
--- a/internal/users/users.go
+++ b/internal/users/users.go
@@ -186,7 +186,7 @@ func (u *Users) watchAPIAuths() {
 
 func (u *Users) closeAllConnections() {
 	for _, user := range u.users {
-		user.closeAllConnections()
+		user.CloseAllConnections()
 	}
 }
 
diff --git a/pkg/pmapi/mocks/mocks.go b/pkg/pmapi/mocks/mocks.go
index ecf7103c..1e37cc2a 100644
--- a/pkg/pmapi/mocks/mocks.go
+++ b/pkg/pmapi/mocks/mocks.go
@@ -5,12 +5,11 @@
 package mocks
 
 import (
-	io "io"
-	reflect "reflect"
-
 	crypto "github.com/ProtonMail/gopenpgp/v2/crypto"
 	pmapi "github.com/ProtonMail/proton-bridge/pkg/pmapi"
 	gomock "github.com/golang/mock/gomock"
+	io "io"
+	reflect "reflect"
 )
 
 // MockClient is a mock of Client interface