Other: Fix flaky cookies test

internal/bridge/bridge.go

@@ -14,7 +14,6 @@ import (
 	"github.com/ProtonMail/gluon"
 	"github.com/ProtonMail/gluon/watcher"
 	"github.com/ProtonMail/proton-bridge/v2/internal/constants"
-	"github.com/ProtonMail/proton-bridge/v2/internal/cookies"
 	"github.com/ProtonMail/proton-bridge/v2/internal/events"
 	"github.com/ProtonMail/proton-bridge/v2/internal/focus"
 	"github.com/ProtonMail/proton-bridge/v2/internal/user"
@@ -35,7 +34,6 @@ type Bridge struct {
 
 	// api manages user API clients.
 	api        *liteapi.Manager
-	cookieJar  *cookies.Jar
 	proxyCtl   ProxyController
 	identifier Identifier
 
@@ -82,24 +80,22 @@ type Bridge struct {
 
 // New creates a new bridge.
 func New(
-	apiURL string, // the URL of the API to use
 	locator Locator, // the locator to provide paths to store data
 	vault *vault.Vault, // the bridge's encrypted data store
+	autostarter Autostarter, // the autostarter to manage autostart settings
+	updater Updater, // the updater to fetch and install updates
+	curVersion *semver.Version, // the current version of the bridge
+
+	apiURL string, // the URL of the API to use
+	cookieJar http.CookieJar, // the cookie jar to use
 	identifier Identifier, // the identifier to keep track of the user agent
 	tlsReporter TLSReporter, // the TLS reporter to report TLS errors
 	roundTripper http.RoundTripper, // the round tripper to use for API requests
 	proxyCtl ProxyController, // the DoH controller
-	autostarter Autostarter, // the autostarter to manage autostart settings
-	updater Updater, // the updater to fetch and install updates
-	curVersion *semver.Version, // the current version of the bridge
+
 	logIMAPClient, logIMAPServer bool, // whether to log IMAP client/server activity
 	logSMTP bool, // whether to log SMTP activity
 ) (*Bridge, error) {
-	cookieJar, err := cookies.NewCookieJar(vault)
-	if err != nil {
-		return nil, fmt.Errorf("failed to create cookie jar: %w", err)
-	}
-
 	api := liteapi.New(
 		liteapi.WithHostURL(apiURL),
 		liteapi.WithAppVersion(constants.AppVersion),
@@ -133,19 +129,20 @@ func New(
 	}
 
 	bridge := newBridge(
+		locator,
 		vault,
+		autostarter,
+		updater,
+		curVersion,
+
 		api,
-		cookieJar,
-		proxyCtl,
 		identifier,
+		proxyCtl,
+
 		tlsConfig,
 		imapServer,
 		smtpBackend,
-		updater,
-		curVersion,
 		focusService,
-		autostarter,
-		locator,
 		logIMAPClient,
 		logIMAPServer,
 		logSMTP,
@@ -159,19 +156,20 @@ func New(
 }
 
 func newBridge(
+	locator Locator,
 	vault *vault.Vault,
+	autostarter Autostarter,
+	updater Updater,
+	curVersion *semver.Version,
+
 	api *liteapi.Manager,
-	cookieJar *cookies.Jar,
-	proxyCtl ProxyController,
 	identifier Identifier,
+	proxyCtl ProxyController,
+
 	tlsConfig *tls.Config,
 	imapServer *gluon.Server,
 	smtpBackend *smtpBackend,
-	updater Updater,
-	curVersion *semver.Version,
 	focusService *focus.Service,
-	autostarter Autostarter,
-	locator Locator,
 	logIMAPClient, logIMAPServer, logSMTP bool,
 ) *Bridge {
 	return &Bridge{
@@ -179,7 +177,6 @@ func newBridge(
 		users: make(map[string]*user.User),
 
 		api:        api,
-		cookieJar:  cookieJar,
 		proxyCtl:   proxyCtl,
 		identifier: identifier,
 
@@ -308,11 +305,6 @@ func (bridge *Bridge) Close(ctx context.Context) error {
 		}
 	}
 
-	// Persist the cookies.
-	if err := bridge.cookieJar.PersistCookies(); err != nil {
-		logrus.WithError(err).Error("Failed to persist cookies")
-	}
-
 	// Close the focus service.
 	bridge.focusService.Close()
 

internal/bridge/bridge_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/ProtonMail/gopenpgp/v2/crypto"
 	"github.com/ProtonMail/proton-bridge/v2/internal/bridge"
 	"github.com/ProtonMail/proton-bridge/v2/internal/certs"
+	"github.com/ProtonMail/proton-bridge/v2/internal/cookies"
 	"github.com/ProtonMail/proton-bridge/v2/internal/events"
 	"github.com/ProtonMail/proton-bridge/v2/internal/focus"
 	"github.com/ProtonMail/proton-bridge/v2/internal/locations"
@@ -128,7 +129,7 @@ func TestBridge_UserAgent(t *testing.T) {
 			require.NoError(t, err)
 
 			// Assert that the user agent was sent to the API.
-			require.Contains(t, calls[len(calls)-1].Header.Get("User-Agent"), bridge.GetCurrentUserAgent())
+			require.Contains(t, calls[len(calls)-1].RequestHeader.Get("User-Agent"), bridge.GetCurrentUserAgent())
 		})
 	})
 }
@@ -137,9 +138,9 @@ func TestBridge_Cookies(t *testing.T) {
 	withTLSEnv(t, func(ctx context.Context, s *server.Server, netCtl *liteapi.NetCtl, locator bridge.Locator, vaultKey []byte) {
 		sessionIDs := safe.NewSet[string]()
 
-		// Save any session IDs the API returns.
+		// Save any session IDs we use.
 		s.AddCallWatcher(func(call server.Call) {
-			cookie, err := (&http.Request{Header: call.Header}).Cookie("Session-Id")
+			cookie, err := (&http.Request{Header: call.RequestHeader}).Cookie("Session-Id")
 			if err != nil {
 				return
 			}
@@ -398,18 +399,29 @@ func withBridge(
 	require.NoError(t, vault.SetIMAPPort(0))
 	require.NoError(t, vault.SetSMTPPort(0))
 
+	// Create a new cookie jar.
+	cookieJar, err := cookies.NewCookieJar(bridge.NewTestCookieJar(), vault)
+	require.NoError(t, err)
+	defer func() { require.NoError(t, cookieJar.PersistCookies()) }()
+
 	// Create a new bridge.
 	bridge, err := bridge.New(
-		apiURL,
+		// The app stuff.
 		locator,
 		vault,
+		mocks.Autostarter,
+		mocks.Updater,
+		v2_3_0,
+
+		// The API stuff.
+		apiURL,
+		cookieJar,
 		useragent.New(),
 		mocks.TLSReporter,
 		liteapi.NewDialer(netCtl, &tls.Config{InsecureSkipVerify: true}).GetRoundTripper(),
 		mocks.ProxyCtl,
-		mocks.Autostarter,
-		mocks.Updater,
-		v2_3_0,
+
+		// The logging stuff.
 		false,
 		false,
 		false,

internal/bridge/errors.go

@@ -3,6 +3,9 @@ package bridge
 import "errors"
 
 var (
+	ErrVaultInsecure = errors.New("the vault is insecure")
+	ErrVaultCorrupt  = errors.New("the vault is corrupt")
+
 	ErrServeIMAP    = errors.New("failed to serve IMAP")
 	ErrServeSMTP    = errors.New("failed to serve SMTP")
 	ErrWatchUpdates = errors.New("failed to watch for updates")

internal/bridge/mocks.go

@@ -1,6 +1,8 @@
 package bridge
 
 import (
+	"net/http"
+	"net/url"
 	"os"
 	"testing"
 
@@ -41,6 +43,24 @@ func (mocks *Mocks) Close() {
 	close(mocks.TLSIssueCh)
 }
 
+type TestCookieJar struct {
+	cookies map[string][]*http.Cookie
+}
+
+func NewTestCookieJar() *TestCookieJar {
+	return &TestCookieJar{
+		cookies: make(map[string][]*http.Cookie),
+	}
+}
+
+func (j *TestCookieJar) SetCookies(u *url.URL, cookies []*http.Cookie) {
+	j.cookies[u.Host] = cookies
+}
+
+func (j *TestCookieJar) Cookies(u *url.URL) []*http.Cookie {
+	return j.cookies[u.Host]
+}
+
 type TestLocationsProvider struct {
 	config, cache string
 }