feat: switch to proxy when need be

pkg/pmapi/clientmanager.go

@@ -1,24 +1,32 @@
 package pmapi
 
 import (
+	"net/http"
 	"sync"
 
 	"github.com/getsentry/raven-go"
+	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 )
 
 // ClientManager is a manager of clients.
 type ClientManager struct {
+	config *ClientConfig
+
 	clients       map[string]*Client
 	clientsLocker sync.Locker
 
 	tokens       map[string]string
 	tokensLocker sync.Locker
 
-	config *ClientConfig
+	url       string
+	urlLocker sync.Locker
 
 	bridgeAuths chan ClientAuth
 	clientAuths chan ClientAuth
+
+	allowProxy    bool
+	proxyProvider *proxyProvider
 }
 
 type ClientAuth struct {
@@ -33,13 +41,21 @@ func NewClientManager(config *ClientConfig) (cm *ClientManager) {
 	}
 
 	cm = &ClientManager{
+		config: config,
+
 		clients:       make(map[string]*Client),
 		clientsLocker: &sync.Mutex{},
-		tokens:        make(map[string]string),
-		tokensLocker:  &sync.Mutex{},
-		config:        config,
-		bridgeAuths:   make(chan ClientAuth),
-		clientAuths:   make(chan ClientAuth),
+
+		tokens:       make(map[string]string),
+		tokensLocker: &sync.Mutex{},
+
+		url:       RootURL,
+		urlLocker: &sync.Mutex{},
+
+		bridgeAuths: make(chan ClientAuth),
+		clientAuths: make(chan ClientAuth),
+
+		proxyProvider: newProxyProvider(dohProviders, proxyQuery),
 	}
 
 	go cm.forwardClientAuths()
@@ -47,6 +63,12 @@ func NewClientManager(config *ClientConfig) (cm *ClientManager) {
 	return
 }
 
+// SetClientRoundTripper sets the roundtripper used by clients created by this client manager.
+func (cm *ClientManager) SetClientRoundTripper(rt http.RoundTripper) {
+	logrus.Info("Setting client roundtripper")
+	cm.config.Transport = rt
+}
+
 // GetClient returns a client for the given userID.
 // If the client does not exist already, it is created.
 func (cm *ClientManager) GetClient(userID string) *Client {
@@ -71,7 +93,7 @@ func (cm *ClientManager) LogoutClient(userID string) {
 
 	go func() {
 		if err := client.logout(); err != nil {
-			// TODO: Try again!
+			// TODO: Try again! This should loop until it succeeds (might fail the first time due to internet).
 			logrus.WithError(err).Error("Client logout failed, not trying again")
 		}
 		client.clearSensitiveData()
@@ -81,6 +103,56 @@ func (cm *ClientManager) LogoutClient(userID string) {
 	return
 }
 
+// GetRootURL returns the root URL to make requests to.
+// It does not include the protocol i.e. no "https://".
+func (cm *ClientManager) GetRootURL() string {
+	cm.urlLocker.Lock()
+	defer cm.urlLocker.Unlock()
+
+	return cm.url
+}
+
+// SetRootURL sets the root URL to make requests to.
+func (cm *ClientManager) SetRootURL(url string) {
+	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.allowProxy = true
+}
+
+// DisallowProxy prevents the client manager from switching clients over to a proxy if need be.
+func (cm *ClientManager) DisallowProxy() {
+	cm.allowProxy = false
+}
+
+// FindProxy returns a usable proxy server.
+func (cm *ClientManager) SwitchToProxy() (proxy string, err error) {
+	logrus.Info("Attempting gto switch to a proxy")
+
+	if proxy, err = cm.proxyProvider.findProxy(); err != nil {
+		err = errors.Wrap(err, "failed to find usable proxy")
+		return
+	}
+
+	cm.SetRootURL(proxy)
+
+	// TODO: Disable after 24 hours.
+
+	return
+}
+
 // GetConfig returns the config used to configure clients.
 func (cm *ClientManager) GetConfig() *ClientConfig {
 	return cm.config
@@ -113,7 +185,7 @@ func (cm *ClientManager) setToken(userID, token string) {
 	cm.tokensLocker.Lock()
 	defer cm.tokensLocker.Unlock()
 
-	logrus.WithField("userID", userID).WithField("token", token).Info("Updating refresh token")
+	logrus.WithField("userID", userID).Info("Updating refresh token")
 
 	cm.tokens[userID] = token
 }
@@ -127,16 +199,19 @@ func (cm *ClientManager) clearToken(userID string) {
 	delete(cm.tokens, userID)
 }
 
-// handleClientAuth
+// handleClientAuth updates or clears client authorisation based on auths received.
 func (cm *ClientManager) handleClientAuth(ca ClientAuth) {
-	// TODO: Maybe want to logout the client in case of nil auth.
+	// If we aren't managing this client, there's nothing to do.
 	if _, ok := cm.clients[ca.UserID]; !ok {
 		return
 	}
 
+	// If the auth is nil, we should clear the token.
+	// TODO: Maybe we should trigger a client logout here? Then we don't have to remember to log it out ourself.
 	if ca.Auth == nil {
 		cm.clearToken(ca.UserID)
-	} else {
-		cm.setToken(ca.UserID, ca.Auth.GenToken())
+		return
 	}
+
+	cm.setToken(ca.UserID, ca.Auth.GenToken())
 }