feat: implement token expiration watcher

internal/bridge/bridge.go

@@ -181,10 +181,11 @@ func (b *Bridge) watchBridgeOutdated() {
 // watchUserAuths receives auths from the client manager and sends them to the appropriate user.
 func (b *Bridge) watchUserAuths() {
 	for auth := range b.clientManager.GetBridgeAuthChannel() {
-		logrus.WithField("token", auth.Auth.GenToken()).WithField("userID", auth.UserID).Info("Received auth from bridge auth channel")
+		logrus.Debug("Bridge received auth from ClientManager")
 
 		if user, ok := b.hasUser(auth.UserID); ok {
-			user.ReceiveAPIAuth(auth.Auth)
+			logrus.Debug("Bridge is forwarding auth to user")
+			user.AuthorizeWithAPIAuth(auth.Auth)
 		} else {
 			logrus.Info("User is not added to bridge yet")
 		}
@@ -494,11 +495,7 @@ func (b *Bridge) updateCurrentUserAgent() {
 
 // hasUser returns whether the bridge currently has a user with ID `id`.
 func (b *Bridge) hasUser(id string) (user *User, ok bool) {
-	logrus.WithField("id", id).Info("Checking whether bridge has given user")
-
 	for _, u := range b.users {
-		logrus.WithField("id", u.ID()).Info("Found potential user")
-
 		if u.ID() == id {
 			user, ok = u, true
 			return

internal/bridge/user.go

@@ -243,10 +243,12 @@ func (u *User) authorizeAndUnlock() (err error) {
 	return nil
 }
 
-func (u *User) ReceiveAPIAuth(auth *pmapi.Auth) {
+func (u *User) AuthorizeWithAPIAuth(auth *pmapi.Auth) {
 	u.lock.Lock()
 	defer u.lock.Unlock()
 
+	u.log.Debug("User received auth from bridge")
+
 	if auth == nil {
 		if err := u.logout(); err != nil {
 			u.log.WithError(err).Error("Failed to logout user after receiving empty auth from API")

pkg/pmapi/auth.go

@@ -24,7 +24,6 @@ import (
 	"fmt"
 	"net/http"
 	"strings"
-	"time"
 
 	pmcrypto "github.com/ProtonMail/gopenpgp/crypto"
 	"github.com/ProtonMail/proton-bridge/pkg/srp"
@@ -197,6 +196,9 @@ type AuthRefreshReq struct {
 }
 
 func (c *Client) sendAuth(auth *Auth) {
+	go func() {
+		c.log.Debug("Client is sending auth to ClientManager")
+
 		c.cm.getClientAuthChannel() <- ClientAuth{
 			UserID: c.userID,
 			Auth:   auth,
@@ -206,6 +208,7 @@ func (c *Client) sendAuth(auth *Auth) {
 			c.uid = auth.UID()
 			c.accessToken = auth.accessToken
 		}
+	}()
 }
 
 // AuthInfo gets authentication info for a user.
@@ -324,7 +327,6 @@ func (c *Client) Auth(username, password string, info *AuthInfo) (auth *Auth, er
 		}
 	}
 
-	c.expiresAt = time.Now().Add(time.Duration(auth.ExpiresIn) * time.Second)
 	return auth, err
 }
 
@@ -445,7 +447,6 @@ func (c *Client) AuthRefresh(uidAndRefreshToken string) (auth *Auth, err error)
 
 	auth = res.getAuth()
 	c.sendAuth(auth)
-	c.expiresAt = time.Now().Add(time.Duration(auth.ExpiresIn) * time.Second)
 
 	return auth, err
 }

pkg/pmapi/client.go

@@ -81,6 +81,7 @@ type ClientConfig struct {
 
 	// Transport specifies the mechanism by which individual HTTP requests are made.
 	// If nil, http.DefaultTransport is used.
+	// TODO: This could be removed entirely and set in the client manager via SetClientRoundTripper.
 	Transport http.RoundTripper
 
 	// Timeout specifies the timeout from request to getting response headers to our API.
@@ -108,7 +109,6 @@ type Client struct {
 	requestLocker sync.Locker
 	keyLocker     sync.Locker
 
-	expiresAt time.Time
 	user      *User
 	addresses AddressList
 	kr        *pmcrypto.KeyRing

pkg/pmapi/clientmanager.go

@@ -3,12 +3,15 @@ package pmapi
 import (
 	"net/http"
 	"sync"
+	"time"
 
 	"github.com/getsentry/raven-go"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 )
 
+var proxyUseDuration = 24 * time.Hour
+
 // ClientManager is a manager of clients.
 type ClientManager struct {
 	config *ClientConfig
@@ -17,6 +20,7 @@ type ClientManager struct {
 	clientsLocker sync.Locker
 
 	tokens           map[string]string
+	tokenExpirations map[string]*tokenExpiration
 	tokensLocker     sync.Locker
 
 	url       string
@@ -34,6 +38,11 @@ type ClientAuth struct {
 	Auth   *Auth
 }
 
+type tokenExpiration struct {
+	timer  *time.Timer
+	cancel chan (struct{})
+}
+
 // NewClientManager creates a new ClientMan which manages clients configured with the given client config.
 func NewClientManager(config *ClientConfig) (cm *ClientManager) {
 	if err := raven.SetDSN(config.SentryDSN); err != nil {
@@ -47,6 +56,7 @@ func NewClientManager(config *ClientConfig) (cm *ClientManager) {
 		clientsLocker: &sync.Mutex{},
 
 		tokens:           make(map[string]string),
+		tokenExpirations: make(map[string]*tokenExpiration),
 		tokensLocker:     &sync.Mutex{},
 
 		url:       RootURL,
@@ -112,43 +122,56 @@ func (cm *ClientManager) GetRootURL() string {
 	return cm.url
 }
 
-// SetRootURL sets the root URL to make requests to.
+// IsProxyAllowed returns whether the user has allowed us to switch to a proxy if need be.
-func (cm *ClientManager) SetRootURL(url string) {
+func (cm *ClientManager) IsProxyAllowed() bool {
 	cm.urlLocker.Lock()
 	defer cm.urlLocker.Unlock()
 
-	logrus.WithField("url", url).Info("Changing to a new root URL")
-
-	cm.url = url
-}
-
-// IsProxyAllowed returns whether the user has allowed us to switch to a proxy if need be.
-func (cm *ClientManager) IsProxyAllowed() bool {
 	return cm.allowProxy
 }
 
 // AllowProxy allows the client manager to switch clients over to a proxy if need be.
 func (cm *ClientManager) AllowProxy() {
+	cm.urlLocker.Lock()
+	defer cm.urlLocker.Unlock()
+
 	cm.allowProxy = true
 }
 
 // DisallowProxy prevents the client manager from switching clients over to a proxy if need be.
 func (cm *ClientManager) DisallowProxy() {
+	cm.urlLocker.Lock()
+	defer cm.urlLocker.Unlock()
+
 	cm.allowProxy = false
+	cm.url = RootURL
+}
+
+// IsProxyEnabled returns whether we are currently proxying requests.
+func (cm *ClientManager) IsProxyEnabled() bool {
+	cm.urlLocker.Lock()
+	defer cm.urlLocker.Unlock()
+
+	return cm.url != RootURL
 }
 
 // FindProxy returns a usable proxy server.
 func (cm *ClientManager) SwitchToProxy() (proxy string, err error) {
-	logrus.Info("Attempting gto switch to a proxy")
+	cm.urlLocker.Lock()
+	defer cm.urlLocker.Unlock()
+
+	logrus.Info("Attempting to switch to a proxy")
 
 	if proxy, err = cm.proxyProvider.findProxy(); err != nil {
-		err = errors.Wrap(err, "failed to find usable proxy")
+		err = errors.Wrap(err, "failed to find a usable proxy")
 		return
 	}
 
-	cm.SetRootURL(proxy)
+	logrus.WithField("proxy", proxy).Info("Switching to a proxy")
 
-	// TODO: Disable after 24 hours.
+	cm.url = proxy
+
+	// TODO: Disable again after 24 hours.
 
 	return
 }
@@ -165,7 +188,7 @@ func (cm *ClientManager) GetToken(userID string) string {
 
 // GetBridgeAuthChannel returns a channel on which client auths can be received.
 func (cm *ClientManager) GetBridgeAuthChannel() chan ClientAuth {
-	return cm.clientAuths
+	return cm.bridgeAuths
 }
 
 // getClientAuthChannel returns a channel on which clients should send auths.
@@ -176,25 +199,46 @@ func (cm *ClientManager) getClientAuthChannel() chan ClientAuth {
 // forwardClientAuths handles all incoming auths from clients before forwarding them on the bridge auth channel.
 func (cm *ClientManager) forwardClientAuths() {
 	for auth := range cm.clientAuths {
+		logrus.Debug("ClientManager received auth from client")
 		cm.handleClientAuth(auth)
+		logrus.Debug("ClientManager is forwarding auth to bridge")
 		cm.bridgeAuths <- auth
 	}
 }
 
-func (cm *ClientManager) setToken(userID, token string) {
+// setToken sets the token for the given userID with the given expiration time.
+func (cm *ClientManager) setToken(userID, token string, expiration time.Duration) {
 	cm.tokensLocker.Lock()
 	defer cm.tokensLocker.Unlock()
 
-	logrus.WithField("userID", userID).Info("Updating refresh token")
+	logrus.WithField("userID", userID).Info("Updating token")
 
 	cm.tokens[userID] = token
+
+	cm.setTokenExpiration(userID, expiration)
+}
+
+// setTokenExpiration will ensure the token is refreshed if it expires.
+// If the token already has an expiration time set, it is replaced.
+func (cm *ClientManager) setTokenExpiration(userID string, expiration time.Duration) {
+	if exp, ok := cm.tokenExpirations[userID]; ok {
+		exp.timer.Stop()
+		close(exp.cancel)
+	}
+
+	cm.tokenExpirations[userID] = &tokenExpiration{
+		timer:  time.NewTimer(expiration),
+		cancel: make(chan struct{}),
+	}
+
+	go cm.watchTokenExpiration(userID)
 }
 
 func (cm *ClientManager) clearToken(userID string) {
 	cm.tokensLocker.Lock()
 	defer cm.tokensLocker.Unlock()
 
-	logrus.WithField("userID", userID).Info("Clearing refresh token")
+	logrus.WithField("userID", userID).Info("Clearing token")
 
 	delete(cm.tokens, userID)
 }
@@ -203,6 +247,7 @@ func (cm *ClientManager) clearToken(userID string) {
 func (cm *ClientManager) handleClientAuth(ca ClientAuth) {
 	// If we aren't managing this client, there's nothing to do.
 	if _, ok := cm.clients[ca.UserID]; !ok {
+		logrus.WithField("userID", ca.UserID).Info("Handling auth for unmanaged client")
 		return
 	}
 
@@ -213,5 +258,18 @@ func (cm *ClientManager) handleClientAuth(ca ClientAuth) {
 		return
 	}
 
-	cm.setToken(ca.UserID, ca.Auth.GenToken())
+	cm.setToken(ca.UserID, ca.Auth.GenToken(), time.Duration(ca.Auth.ExpiresIn)*time.Second)
+}
+
+func (cm *ClientManager) watchTokenExpiration(userID string) {
+	expiration := cm.tokenExpirations[userID]
+
+	select {
+	case <-expiration.timer.C:
+		logrus.WithField("userID", userID).Info("Auth token expired! Refreshing")
+		cm.clients[userID].AuthRefresh(cm.tokens[userID])
+
+	case <-expiration.cancel:
+		logrus.WithField("userID", userID).Info("Auth was refreshed before it expired, cancelling this watcher")
+	}
 }

pkg/pmapi/dialer_with_proxy.go

@@ -282,13 +282,15 @@ func (p *DialerWithPinning) dialAndCheckFingerprints(network, address string) (c
 func (p *DialerWithPinning) dialWithProxyFallback(network, address string) (conn net.Conn, err error) {
 	p.log.Info("Dialing with proxy fallback")
 
-	var host, port string
+	// Try to dial, and if it succeeds, then just return.
-	if host, port, err = net.SplitHostPort(address); err != nil {
+	if conn, err = p.dial(network, address); err == nil {
 		return
 	}
 
-	// Try to dial, and if it succeeds, then just return.
+	p.log.WithField("address", address).WithError(err).Error("Dialing failed")
-	if conn, err = p.dial(network, address); err == nil {
+
+	host, port, err := net.SplitHostPort(address)
+	if err != nil {
 		return
 	}
 
@@ -296,18 +298,21 @@ func (p *DialerWithPinning) dialWithProxyFallback(network, address string) (conn
 	// (e.g. we dial protonmail.com/... to check for updates), there's also no point in
 	// continuing since a proxy won't help us reach that.
 	if !p.cm.IsProxyAllowed() || host != p.cm.GetRootURL() {
-		p.log.WithField("useProxy", p.cm.IsProxyAllowed()).Info("Dial failed but not switching to proxy")
+		p.log.WithField("address", address).Debug("Aborting dial, cannot switch to a proxy")
 		return
 	}
 
 	// Switch to a proxy and retry the dial.
-	var proxy string
+	proxy, err := p.cm.SwitchToProxy()
-
+	if err != nil {
-	if proxy, err = p.cm.SwitchToProxy(); err != nil {
 		return
 	}
 
-	return p.dial(network, net.JoinHostPort(proxy, port))
+	proxyAddress := net.JoinHostPort(proxy, port)
+
+	p.log.WithField("address", proxyAddress).Debug("Trying dial again using a proxy")
+
+	return p.dial(network, proxyAddress)
 }
 
 // dial returns a connection to the given address using the given network.

pkg/pmapi/proxy.go

@@ -51,7 +51,6 @@ type proxyProvider struct {
 	query      string   // The query string used to find proxies.
 	proxyCache []string // All known proxies, cached in case DoH providers are unreachable.
 
-	useDuration                time.Duration // How much time to use the proxy before returning to the original API.
 	findTimeout, lookupTimeout time.Duration // Timeouts for DNS query and proxy search.
 
 	lastLookup time.Time // The time at which we last attempted to find a proxy.
@@ -63,7 +62,6 @@ func newProxyProvider(providers []string, query string) (p *proxyProvider) { //
 	p = &proxyProvider{
 		providers:     providers,
 		query:         query,
-		useDuration:   proxyRevertTime,
 		findTimeout:   proxySearchTimeout,
 		lookupTimeout: proxyQueryTimeout,
 	}
@@ -148,6 +146,7 @@ func (p *proxyProvider) canReach(url string) bool {
 		SetTLSClientConfig(&tls.Config{InsecureSkipVerify: true}) // nolint[gosec]
 
 	if _, err := pinger.R().Get("/tests/ping"); err != nil {
+		logrus.WithField("proxy", url).WithError(err).Warn("Failed to ping proxy")
 		return false
 	}
 

pkg/pmapi/users.go

@@ -109,7 +109,6 @@ func (c *Client) UpdateUser() (user *User, err error) {
 	}
 
 	c.user = user
-	c.log.Infoln("update user:", user.ID)
 	raven.SetUserContext(&raven.User{ID: user.ID})
 
 	var tmpList AddressList
@@ -117,6 +116,8 @@ func (c *Client) UpdateUser() (user *User, err error) {
 		c.addresses = tmpList
 	}
 
+	c.log.WithField("userID", user.ID).Info("Updated user")
+
 	return user, err
 }
 

test/context/bridge.go

@@ -24,6 +24,7 @@ import (
 	"github.com/ProtonMail/proton-bridge/internal/bridge"
 	"github.com/ProtonMail/proton-bridge/internal/preferences"
 	"github.com/ProtonMail/proton-bridge/pkg/listener"
+	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
 )
 
 // GetBridge returns bridge instance.
@@ -34,10 +35,7 @@ func (ctx *TestContext) GetBridge() *bridge.Bridge {
 // withBridgeInstance creates a bridge instance for use in the test.
 // Every TestContext has this by default and thus this doesn't need to be exported.
 func (ctx *TestContext) withBridgeInstance() {
-	pmapiFactory := func(userID string) bridge.PMAPIProvider {
+	ctx.bridge = newBridgeInstance(ctx.t, ctx.cfg, ctx.credStore, ctx.listener, ctx.clientManager)
-		return ctx.pmapiController.GetClient(userID)
-	}
-	ctx.bridge = newBridgeInstance(ctx.t, ctx.cfg, ctx.credStore, ctx.listener, pmapiFactory)
 	ctx.addCleanupChecked(ctx.bridge.ClearData, "Cleaning bridge data")
 }
 
@@ -62,7 +60,7 @@ func newBridgeInstance(
 	cfg *fakeConfig,
 	credStore bridge.CredentialsStorer,
 	eventListener listener.Listener,
-	pmapiFactory bridge.PMAPIProviderFactory,
+	clientManager *pmapi.ClientManager,
 ) *bridge.Bridge {
 	version := os.Getenv("VERSION")
 	bridge.UpdateCurrentUserAgent(version, runtime.GOOS, "", "")
@@ -70,7 +68,7 @@ func newBridgeInstance(
 	panicHandler := &panicHandler{t: t}
 	pref := preferences.New(cfg)
 
-	return bridge.New(cfg, pref, panicHandler, eventListener, version, pmapiFactory, credStore)
+	return bridge.New(cfg, pref, panicHandler, eventListener, version, clientManager, credStore)
 }
 
 // SetLastBridgeError sets the last error that occurred while executing a bridge action.

test/context/config.go

@@ -28,7 +28,6 @@ import (
 
 type fakeConfig struct {
 	dir string
-	tm  *pmapi.TokenManager
 }
 
 // newFakeConfig creates a temporary folder for files.
@@ -41,7 +40,6 @@ func newFakeConfig() *fakeConfig {
 
 	return &fakeConfig{
 		dir: dir,
-		tm:  pmapi.NewTokenManager(),
 	}
 }
 
@@ -52,8 +50,7 @@ func (c *fakeConfig) GetAPIConfig() *pmapi.ClientConfig {
 	return &pmapi.ClientConfig{
 		AppVersion: "Bridge_" + os.Getenv("VERSION"),
 		ClientID:   "bridge",
-		// TokenManager should not be required, but PMAPI still doesn't handle not-set cases everywhere.
+		SentryDSN:  "",
-		TokenManager: c.tm,
 	}
 }
 func (c *fakeConfig) GetDBDir() string {

test/context/context.go

@@ -21,6 +21,7 @@ package context
 import (
 	"github.com/ProtonMail/proton-bridge/internal/bridge"
 	"github.com/ProtonMail/proton-bridge/pkg/listener"
+	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
 	"github.com/ProtonMail/proton-bridge/test/accounts"
 	"github.com/ProtonMail/proton-bridge/test/mocks"
 	"github.com/sirupsen/logrus"
@@ -57,6 +58,9 @@ type TestContext struct {
 	smtpClients       map[string]*mocks.SMTPClient
 	smtpLastResponses map[string]*mocks.SMTPResponse
 
+	// PMAPI related variables.
+	clientManager *pmapi.ClientManager
+
 	// These are the cleanup steps executed when Cleanup() is called.
 	cleanupSteps []*Cleaner
 
@@ -70,17 +74,20 @@ func New() *TestContext {
 
 	cfg := newFakeConfig()
 
+	cm := pmapi.NewClientManager(cfg.GetAPIConfig())
+
 	ctx := &TestContext{
 		t:                 &bddT{},
 		cfg:               cfg,
 		listener:          listener.New(),
-		pmapiController:   newPMAPIController(),
+		pmapiController:   newPMAPIController(cm),
 		testAccounts:      newTestAccounts(),
 		credStore:         newFakeCredStore(),
 		imapClients:       make(map[string]*mocks.IMAPClient),
 		imapLastResponses: make(map[string]*mocks.IMAPResponse),
 		smtpClients:       make(map[string]*mocks.SMTPClient),
 		smtpLastResponses: make(map[string]*mocks.SMTPResponse),
+		clientManager:     cm,
 		logger:            logrus.StandardLogger(),
 	}
 

test/context/pmapi_controller.go

@@ -40,12 +40,12 @@ type PMAPIController interface {
 	GetCalls(method, path string) [][]byte
 }
 
-func newPMAPIController() PMAPIController {
+func newPMAPIController(cm *pmapi.ClientManager) PMAPIController {
 	switch os.Getenv(EnvName) {
 	case EnvFake:
 		return newFakePMAPIController()
 	case EnvLive:
-		return newLivePMAPIController()
+		return newLivePMAPIController(cm)
 	default:
 		panic("unknown env")
 	}
@@ -67,8 +67,8 @@ func (s *fakePMAPIControllerWrap) GetClient(userID string) bridge.PMAPIProvider
 	return s.Controller.GetClient(userID)
 }
 
-func newLivePMAPIController() PMAPIController {
+func newLivePMAPIController(cm *pmapi.ClientManager) PMAPIController {
-	return newLiveAPIControllerWrap(liveapi.NewController())
+	return newLiveAPIControllerWrap(liveapi.NewController(cm))
 }
 
 type liveAPIControllerWrap struct {

test/fakeapi/auth.go

@@ -141,13 +141,12 @@ func (api *FakePMAPI) AuthRefresh(token string) (*pmapi.Auth, error) {
 	return auth, nil
 }
 
-func (api *FakePMAPI) Logout() error {
+func (api *FakePMAPI) Logout() {
 	if err := api.checkAndRecordCall(DELETE, "/auth", nil); err != nil {
-		return err
+		return
 	}
 	// Logout will also emit change to auth channel
 	api.sendAuth(nil)
 	api.controller.deleteSession(api.uid)
 	api.unsetUser()
-	return nil
 }

test/liveapi/controller.go

@@ -18,9 +18,7 @@
 package liveapi
 
 import (
-	"fmt"
 	"net/http"
-	"os"
 	"sync"
 
 	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
@@ -32,31 +30,31 @@ type Controller struct {
 	calls                []*fakeCall
 	pmapiByUsername      map[string]*pmapi.Client
 	messageIDsByUsername map[string][]string
+	clientManager        *pmapi.ClientManager
 
 	// State controlled by test.
 	noInternetConnection bool
 }
 
-func NewController() *Controller {
+func NewController(cm *pmapi.ClientManager) *Controller {
-	return &Controller{
+	cntrl := &Controller{
 		lock:                 &sync.RWMutex{},
 		calls:                []*fakeCall{},
 		pmapiByUsername:      map[string]*pmapi.Client{},
 		messageIDsByUsername: map[string][]string{},
+		clientManager:        cm,
 
 		noInternetConnection: false,
 	}
+
+	cntrl.clientManager.SetClientRoundTripper(&fakeTransport{
+		cntrl:     cntrl,
+		transport: http.DefaultTransport,
+	})
+
+	return cntrl
 }
 
 func (cntrl *Controller) GetClient(userID string) *pmapi.Client {
-	cfg := &pmapi.ClientConfig{
+	return cntrl.clientManager.GetClient(userID)
-		AppVersion: fmt.Sprintf("Bridge_%s", os.Getenv("VERSION")),
-		ClientID:   "bridge-test",
-		Transport: &fakeTransport{
-			cntrl:     cntrl,
-			transport: http.DefaultTransport,
-		},
-		TokenManager: pmapi.NewTokenManager(),
-	}
-	return pmapi.NewClient(cfg, userID)
 }

test/liveapi/users.go

@@ -18,9 +18,6 @@
 package liveapi
 
 import (
-	"fmt"
-	"os"
-
 	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
 	"github.com/cucumber/godog"
 	"github.com/pkg/errors"
@@ -31,11 +28,7 @@ func (cntrl *Controller) AddUser(user *pmapi.User, addresses *pmapi.AddressList,
 		return godog.ErrPending
 	}
 
-	client := pmapi.NewClient(&pmapi.ClientConfig{
+	client := cntrl.GetClient(user.ID)
-		AppVersion:   fmt.Sprintf("Bridge_%s", os.Getenv("VERSION")),
-		ClientID:     "bridge-cntrl",
-		TokenManager: pmapi.NewTokenManager(),
-	}, user.ID)
 
 	authInfo, err := client.AuthInfo(user.Name)
 	if err != nil {
@@ -62,5 +55,6 @@ func (cntrl *Controller) AddUser(user *pmapi.User, addresses *pmapi.AddressList,
 	}
 
 	cntrl.pmapiByUsername[user.Name] = client
+
 	return nil
 }