Tests and final touches

internal/imap/cache/cache.go

@@ -23,7 +23,7 @@ import (
 	"sync"
 	"time"
 
-	backendMessage "github.com/ProtonMail/proton-bridge/pkg/message"
+	pkgMsg "github.com/ProtonMail/proton-bridge/pkg/message"
 )
 
 type key struct {
@@ -41,7 +41,7 @@ func (s oldestFirst) Less(i, j int) bool { return s[i].Timestamp < s[j].Timestam
 type cachedMessage struct {
 	key
 	data      []byte
-	structure backendMessage.BodyStructure
+	structure pkgMsg.BodyStructure
 }
 
 //nolint[gochecknoglobals]
@@ -101,7 +101,7 @@ func BuildUnlock(messageID string) {
 	delete(buildLocks, messageID)
 }
 
-func LoadMail(mID string) (reader *bytes.Reader, structure *backendMessage.BodyStructure) {
+func LoadMail(mID string) (reader *bytes.Reader, structure *pkgMsg.BodyStructure) {
 	reader = &bytes.Reader{}
 	cacheMutex.Lock()
 	defer cacheMutex.Unlock()
@@ -115,7 +115,7 @@ func LoadMail(mID string) (reader *bytes.Reader, structure *backendMessage.BodyS
 	return
 }
 
-func SaveMail(mID string, msg []byte, structure *backendMessage.BodyStructure) {
+func SaveMail(mID string, msg []byte, structure *pkgMsg.BodyStructure) {
 	cacheMutex.Lock()
 	defer cacheMutex.Unlock()
 

internal/imap/cache/cache_test.go

@@ -22,11 +22,11 @@ import (
 	"testing"
 	"time"
 
-	bckMsg "github.com/ProtonMail/proton-bridge/pkg/message"
+	pkgMsg "github.com/ProtonMail/proton-bridge/pkg/message"
 	"github.com/stretchr/testify/require"
 )
 
-var bs = &bckMsg.BodyStructure{} //nolint[gochecknoglobals]
+var bs = &pkgMsg.BodyStructure{} //nolint[gochecknoglobals]
 const testUID = "testmsg"
 
 func TestSaveAndLoad(t *testing.T) {

internal/imap/mailbox.go

@@ -197,8 +197,8 @@ func (im *imapMailbox) Expunge() error {
 // UIDExpunge permanently removes messages that have the \Deleted flag set
 // and UID passed from SeqSet from the currently selected mailbox.
 func (im *imapMailbox) UIDExpunge(seqSet *imap.SeqSet) error {
-	im.user.backend.setUpdatesBeBlocking(im.user.currentAddressLowercase, im.name, operationDeleteMessage)
-	defer im.user.backend.unsetUpdatesBeBlocking(im.user.currentAddressLowercase, im.name, operationDeleteMessage)
+	im.user.backend.updates.block(im.user.currentAddressLowercase, im.name, operationDeleteMessage)
+	defer im.user.backend.updates.unblock(im.user.currentAddressLowercase, im.name, operationDeleteMessage)
 
 	messageIDs, err := im.apiIDsFromSeqSet(true, seqSet)
 	if err != nil || len(messageIDs) == 0 {

internal/imap/mailbox_message.go

@@ -328,12 +328,21 @@ func (im *imapMailbox) getBodyAndStructure(storeMessage storeMessageProvider) (
 		var body []byte
 		structure, body, err = im.buildMessage(m)
 		m.Size = int64(len(body))
+		// Save size and body structure even for messages unable to decrypt
+		// so the size or body structure doesn't have to be computed every time.
 		if err := storeMessage.SetSize(m.Size); err != nil {
 			im.log.WithError(err).
 				WithField("newSize", m.Size).
 				WithField("msgID", m.ID).
 				Warn("Cannot update size while building")
 		}
+		if structure != nil && !isMessageInDraftFolder(m) {
+			if err := storeMessage.SetBodyStructure(structure); err != nil {
+				im.log.WithError(err).
+					WithField("msgID", m.ID).
+					Warn("Cannot update bodystructure while building")
+			}
+		}
 		if err == nil && structure != nil && len(body) > 0 {
 			if err := storeMessage.SetContentTypeAndHeader(m.MIMEType, m.Header); err != nil {
 				im.log.WithError(err).
@@ -342,11 +351,6 @@ func (im *imapMailbox) getBodyAndStructure(storeMessage storeMessageProvider) (
 			}
 			// Drafts can change and we don't want to cache them.
 			if !isMessageInDraftFolder(m) {
-				if err := storeMessage.SetBodyStructure(structure); err != nil {
-					im.log.WithError(err).
-						WithField("msgID", m.ID).
-						Warn("Cannot update bodystructure while building")
-				}
 				cache.SaveMail(id, body, structure)
 			}
 			bodyReader = bytes.NewReader(body)

internal/imap/server.go

@@ -23,6 +23,7 @@ import (
 	"io"
 	"net"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	imapid "github.com/ProtonMail/go-imap-id"
@@ -31,6 +32,7 @@ import (
 	"github.com/ProtonMail/proton-bridge/internal/imap/id"
 	"github.com/ProtonMail/proton-bridge/internal/imap/uidplus"
 	"github.com/ProtonMail/proton-bridge/pkg/listener"
+	"github.com/ProtonMail/proton-bridge/pkg/ports"
 	"github.com/emersion/go-imap"
 	imapappendlimit "github.com/emersion/go-imap-appendlimit"
 	imapidle "github.com/emersion/go-imap-idle"
@@ -48,6 +50,8 @@ type imapServer struct {
 	eventListener listener.Listener
 	debugClient   bool
 	debugServer   bool
+	port          int
+	isRunning     atomic.Value
 }
 
 // NewIMAPServer constructs a new IMAP server configured with the given options.
@@ -96,13 +100,16 @@ func NewIMAPServer(panicHandler panicHandler, debugClient, debugServer bool, por
 		uidplus.NewExtension(),
 	)
 
-	return &imapServer{
+	server := &imapServer{
 		panicHandler:  panicHandler,
 		server:        s,
 		eventListener: eventListener,
 		debugClient:   debugClient,
 		debugServer:   debugServer,
+		port:          port,
 	}
+	server.isRunning.Store(false)
+	return server
 }
 
 // Starts the server.
@@ -114,6 +121,11 @@ func (s *imapServer) ListenAndServe() {
 }
 
 func (s *imapServer) listenAndServe() {
+	if s.isRunning.Load().(bool) {
+		return
+	}
+	s.isRunning.Store(true)
+
 	log.Info("IMAP server listening at ", s.server.Addr)
 	l, err := net.Listen("tcp", s.server.Addr)
 	if err != nil {
@@ -126,19 +138,13 @@ func (s *imapServer) listenAndServe() {
 		Listener: l,
 		server:   s,
 	})
-	if err != nil {
-		failed := true
-		if netErr, ok := err.(*net.OpError); ok {
-			originalErr := netErr.Unwrap()
-			if originalErr != nil && originalErr.Error() == "use of closed network connection" {
-				failed = false
-			}
-		}
-		if failed {
-			s.eventListener.Emit(events.ErrorEvent, "IMAP failed: "+err.Error())
-			log.Error("IMAP failed: ", err)
-			return
-		}
+	// Serve returns error every time, even after closing the server.
+	// User shouldn't be notified about error if server shouldn't be running,
+	// but it should in case it was not closed by `s.Close()`.
+	if err != nil && s.isRunning.Load().(bool) {
+		s.eventListener.Emit(events.ErrorEvent, "IMAP failed: "+err.Error())
+		log.Error("IMAP failed: ", err)
+		return
 	}
 	defer s.server.Close() //nolint[errcheck]
 
@@ -147,6 +153,11 @@ func (s *imapServer) listenAndServe() {
 
 // Stops the server.
 func (s *imapServer) Close() {
+	if !s.isRunning.Load().(bool) {
+		return
+	}
+	s.isRunning.Store(false)
+
 	log.Info("Closing IMAP server")
 	if err := s.server.Close(); err != nil {
 		log.WithError(err).Error("Failed to close the connection")
@@ -159,29 +170,32 @@ func (s *imapServer) monitorInternetConnection() {
 	off := make(chan string)
 	s.eventListener.Add(events.InternetOffEvent, off)
 
-	isOn := true
 	for {
+		var expectedIsPortFree bool
 		select {
 		case <-on:
-			if isOn {
-				continue
-			}
-			isOn = true
 			go func() {
 				defer s.panicHandler.HandlePanic()
 				s.listenAndServe()
 			}()
+			expectedIsPortFree = false
 		case <-off:
-			if !isOn {
-				continue
-			}
-			isOn = false
 			s.Close()
+			expectedIsPortFree = true
+		}
+
+		start := time.Now()
+		for {
+			if ports.IsPortFree(s.port) == expectedIsPortFree {
+				break
+			}
+			// Safety stop if something went wrong.
+			if time.Since(start) > 15*time.Second {
+				log.WithField("expectedIsPortFree", expectedIsPortFree).Warn("Server start/stop check timeouted")
+				break
+			}
+			time.Sleep(100 * time.Millisecond)
 		}
-		// Give it some time to serve or close server before changing it again.
-		// E.g., if we get quickly off-on signal, starting or closing could
-		// fail because server is still running or not yet, respectively.
-		time.Sleep(10 * time.Second)
 	}
 }
 

internal/imap/server_test.go

@@ -0,0 +1,65 @@
+// Copyright (c) 2021 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 imap
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/ProtonMail/proton-bridge/internal/bridge"
+	"github.com/ProtonMail/proton-bridge/internal/events"
+	"github.com/ProtonMail/proton-bridge/pkg/listener"
+	"github.com/ProtonMail/proton-bridge/pkg/ports"
+	imapserver "github.com/emersion/go-imap/server"
+
+	"github.com/stretchr/testify/require"
+)
+
+type testPanicHandler struct{}
+
+func (ph *testPanicHandler) HandlePanic() {}
+
+func TestIMAPServerTurnOffAndOnAgain(t *testing.T) {
+	panicHandler := &testPanicHandler{}
+
+	eventListener := listener.New()
+
+	port := ports.FindFreePortFrom(12345)
+	server := imapserver.New(nil)
+	server.Addr = fmt.Sprintf("%v:%v", bridge.Host, port)
+
+	s := &imapServer{
+		panicHandler:  panicHandler,
+		server:        server,
+		eventListener: eventListener,
+	}
+	s.isRunning.Store(false)
+
+	go s.ListenAndServe()
+	time.Sleep(5 * time.Second)
+	require.False(t, ports.IsPortFree(port))
+
+	eventListener.Emit(events.InternetOffEvent, "")
+	time.Sleep(10 * time.Second)
+	require.True(t, ports.IsPortFree(port))
+
+	eventListener.Emit(events.InternetOnEvent, "")
+	time.Sleep(10 * time.Second)
+	require.False(t, ports.IsPortFree(port))
+}

internal/imap/store.go

@@ -24,7 +24,7 @@ import (
 	"github.com/ProtonMail/gopenpgp/v2/crypto"
 	"github.com/ProtonMail/proton-bridge/internal/imap/uidplus"
 	"github.com/ProtonMail/proton-bridge/internal/store"
-	backendMessage "github.com/ProtonMail/proton-bridge/pkg/message"
+	pkgMsg "github.com/ProtonMail/proton-bridge/pkg/message"
 	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
 )
 
@@ -100,8 +100,8 @@ type storeMessageProvider interface {
 
 	SetSize(int64) error
 	SetContentTypeAndHeader(string, mail.Header) error
-	SetBodyStructure(*backendMessage.BodyStructure) error
-	GetBodyStructure() (*backendMessage.BodyStructure, error)
+	SetBodyStructure(*pkgMsg.BodyStructure) error
+	GetBodyStructure() (*pkgMsg.BodyStructure, error)
 }
 
 type storeUserWrap struct {

internal/imap/updates_test.go

@@ -0,0 +1,60 @@
+// Copyright (c) 2021 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 imap
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestUpdatesCanDelete(t *testing.T) {
+	u := newIMAPUpdates()
+
+	can, _ := u.CanDelete("mbox")
+	require.True(t, can)
+
+	u.forbidExpunge("mbox")
+	u.allowExpunge("mbox")
+
+	can, _ = u.CanDelete("mbox")
+	require.True(t, can)
+}
+
+func TestUpdatesCannotDelete(t *testing.T) {
+	u := newIMAPUpdates()
+
+	u.forbidExpunge("mbox")
+	can, wait := u.CanDelete("mbox")
+	require.False(t, can)
+
+	ch := make(chan time.Duration)
+	go func() {
+		start := time.Now()
+		wait()
+		ch <- time.Since(start)
+		close(ch)
+	}()
+
+	time.Sleep(200 * time.Millisecond)
+	u.allowExpunge("mbox")
+	duration := <-ch
+
+	require.True(t, duration > 200*time.Millisecond)
+}