Wrap imap-id with backend caller

internal/imap/currentclient/extension.go

@@ -1,102 +0,0 @@
-// 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 currentclient implements setting client's ID to backend.
-package currentclient
-
-import (
-	"sync"
-	"time"
-
-	imapid "github.com/ProtonMail/go-imap-id"
-	imapserver "github.com/emersion/go-imap/server"
-	"github.com/sirupsen/logrus"
-)
-
-var log = logrus.WithField("pkg", "imap/currentclient") //nolint[gochecknoglobals]
-
-type currentClientSetter interface {
-	SetCurrentClient(name, version string)
-}
-
-type extension struct {
-	backend currentClientSetter
-
-	lastID       imapid.ID
-	lastIDLocker sync.Locker
-}
-
-// NewExtension prepares extension reading IMAP ID from connection
-// and setting it to backend. This is not standard IMAP extension.
-func NewExtension(backend currentClientSetter) imapserver.Extension {
-	return &extension{
-		backend: backend,
-
-		lastID:       imapid.ID{imapid.FieldName: ""},
-		lastIDLocker: &sync.Mutex{},
-	}
-}
-
-func (ext *extension) Capabilities(conn imapserver.Conn) []string {
-	ext.readID(conn)
-	return nil
-}
-
-func (ext *extension) Command(name string) imapserver.HandlerFactory {
-	return nil
-}
-
-func (ext *extension) readID(conn imapserver.Conn) {
-	conn.Server().ForEachConn(func(candidate imapserver.Conn) {
-		if id, ok := candidate.(imapid.Conn); ok {
-			if conn.Context() == candidate.Context() {
-				// ID is not available right at the beginning of the connection.
-				// Clients send ID quickly after AUTH. We need to wait for it.
-				go func() {
-					start := time.Now()
-					for {
-						if id.ID() != nil {
-							ext.setLastID(id.ID())
-							break
-						}
-						if time.Since(start) > 10*time.Second {
-							break
-						}
-						time.Sleep(100 * time.Millisecond)
-					}
-				}()
-			}
-		}
-	})
-}
-
-func (ext *extension) setLastID(id imapid.ID) {
-	ext.lastIDLocker.Lock()
-	defer ext.lastIDLocker.Unlock()
-
-	if name, ok := id[imapid.FieldName]; ok && ext.lastID[imapid.FieldName] != name {
-		ext.lastID = imapid.ID{}
-		for k, v := range id {
-			ext.lastID[k] = v
-		}
-		log.Warn("Mail Client ID changed to ", ext.lastID)
-		ext.backend.SetCurrentClient(
-			ext.lastID[imapid.FieldName],
-			ext.lastID[imapid.FieldVersion],
-		)
-	}
-}

internal/imap/id/extension.go

@@ -0,0 +1,90 @@
+// 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 id
+
+import (
+	imapid "github.com/ProtonMail/go-imap-id"
+	imapserver "github.com/emersion/go-imap/server"
+)
+
+type currentClientSetter interface {
+	SetCurrentClient(name, version string)
+}
+
+// Extension for IMAP server
+type extension struct {
+	extID  imapserver.ConnExtension
+	setter currentClientSetter
+}
+
+func (ext *extension) Capabilities(conn imapserver.Conn) []string {
+	return ext.extID.Capabilities(conn)
+}
+
+func (ext *extension) Command(name string) imapserver.HandlerFactory {
+	newIDHandler := ext.extID.Command(name)
+	if newIDHandler == nil {
+		return nil
+	}
+	return func() imapserver.Handler {
+		if hdlrID, ok := newIDHandler().(*imapid.Handler); ok {
+			return &handler{
+				hdlrID: hdlrID,
+				setter: ext.setter,
+			}
+		}
+		return nil
+	}
+}
+
+func (ext *extension) NewConn(conn imapserver.Conn) imapserver.Conn {
+	return ext.extID.NewConn(conn)
+}
+
+type handler struct {
+	hdlrID *imapid.Handler
+	setter currentClientSetter
+}
+
+func (hdlr *handler) Parse(fields []interface{}) error {
+	return hdlr.hdlrID.Parse(fields)
+}
+
+func (hdlr *handler) Handle(conn imapserver.Conn) error {
+	err := hdlr.hdlrID.Handle(conn)
+	if err == nil {
+		id := hdlr.hdlrID.Command.ID
+		hdlr.setter.SetCurrentClient(
+			id[imapid.FieldName],
+			id[imapid.FieldVersion],
+		)
+	}
+	return err
+}
+
+// NewExtension returns extension which is adding RFC2871 ID capability, with
+// direct interface to set information about email client to backend.
+func NewExtension(serverID imapid.ID, setter currentClientSetter) imapserver.Extension {
+	if conExtID, ok := imapid.NewExtension(serverID).(imapserver.ConnExtension); ok {
+		return &extension{
+			extID:  conExtID,
+			setter: setter,
+		}
+	}
+	return nil
+}

internal/imap/server.go

@@ -28,7 +28,7 @@ import (
 	imapid "github.com/ProtonMail/go-imap-id"
 	"github.com/ProtonMail/proton-bridge/internal/bridge"
 	"github.com/ProtonMail/proton-bridge/internal/events"
-	"github.com/ProtonMail/proton-bridge/internal/imap/currentclient"
+	"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/emersion/go-imap"
@@ -61,7 +61,7 @@ func NewIMAPServer(debugClient, debugServer bool, port int, tls *tls.Config, ima
 	s.UpgradeError = imapBackend.upgradeError
 
 	serverID := imapid.ID{
-		imapid.FieldName:       "ProtonMail",
+		imapid.FieldName:       "ProtonMail Bridge",
 		imapid.FieldVendor:     "Proton Technologies AG",
 		imapid.FieldSupportURL: "https://protonmail.com/support",
 	}
@@ -84,11 +84,10 @@ func NewIMAPServer(debugClient, debugServer bool, port int, tls *tls.Config, ima
 		imapidle.NewExtension(),
 		imapmove.NewExtension(),
 		imapspecialuse.NewExtension(),
-		imapid.NewExtension(serverID),
+		id.NewExtension(serverID, imapBackend.bridge),
 		imapquota.NewExtension(),
 		imapappendlimit.NewExtension(),
 		imapunselect.NewExtension(),
-		currentclient.NewExtension(imapBackend.bridge),
 		uidplus.NewExtension(),
 	)