GODT-35: Finish all details and make tests pass

pkg/pmapi/manager.go

@@ -1,3 +1,20 @@
+// 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 pmapi
 
 import (
@@ -10,56 +27,53 @@ import (
 )
 
 type manager struct {
-	rc *resty.Client
+	cfg Config
+	rc  *resty.Client
 
-	isDown    bool
-	locker    sync.Locker
-	observers []ConnectionObserver
-}
-
-func newManager(cfg Config) *manager {
-	m := &manager{
-		rc:     resty.New(),
-		locker: &sync.Mutex{},
-	}
-
-	// Set the API host.
-	m.rc.SetHostURL(cfg.HostURL)
-
-	// Set static header values.
-	m.rc.SetHeader("x-pm-appversion", cfg.AppVersion)
-
-	// Set middleware.
-	m.rc.OnAfterResponse(catchAPIError)
-
-	// Configure retry mechanism.
-	m.rc.SetRetryMaxWaitTime(time.Minute)
-	m.rc.SetRetryAfter(catchRetryAfter)
-	m.rc.AddRetryCondition(catchTooManyRequests)
-	m.rc.AddRetryCondition(catchNoResponse)
-	m.rc.AddRetryCondition(catchProxyAvailable)
-
-	// Determine what happens when requests succeed/fail.
-	m.rc.OnAfterResponse(m.handleRequestSuccess)
-	m.rc.OnError(m.handleRequestFailure)
-
-	// Set the data type of API errors.
-	m.rc.SetError(&Error{})
-
-	return m
+	isDown              bool
+	locker              sync.Locker
+	connectionObservers []ConnectionObserver
+	proxyDialer         *ProxyTLSDialer
 }
 
 func New(cfg Config) Manager {
 	return newManager(cfg)
 }
 
-func (m *manager) SetLogger(logger resty.Logger) {
-	m.rc.SetLogger(logger)
-	m.rc.SetDebug(true)
+func newManager(cfg Config) *manager {
+	m := &manager{
+		cfg:    cfg,
+		rc:     resty.New(),
+		locker: &sync.Mutex{},
+	}
+
+	proxyDialer, transport := newProxyDialerAndTransport(cfg)
+	m.proxyDialer = proxyDialer
+	m.rc.SetTransport(transport)
+
+	m.rc.SetHostURL(cfg.HostURL)
+	m.rc.OnBeforeRequest(m.setHeaderValues)
+
+	// Any HTTP status code higher than 399 with JSON inside (and proper header)
+	// is converted to Error. `catchAPIError` then processes API custom errors
+	// wrapped in JSON. If error is returned, `handleRequestFailure` is called,
+	// otherwise `handleRequestSuccess` is called.
+	m.rc.SetError(&Error{})
+	m.rc.OnAfterResponse(m.catchAPIError)
+	m.rc.OnAfterResponse(m.handleRequestSuccess)
+	m.rc.OnError(m.handleRequestFailure)
+
+	// Configure retry mechanism.
+	m.rc.SetRetryMaxWaitTime(time.Minute)
+	m.rc.SetRetryAfter(catchRetryAfter)
+	m.rc.AddRetryCondition(shouldRetry)
+
+	return m
 }
 
 func (m *manager) SetTransport(transport http.RoundTripper) {
 	m.rc.SetTransport(transport)
+	m.proxyDialer = nil
 }
 
 func (m *manager) SetCookieJar(jar http.CookieJar) {
@@ -71,7 +85,15 @@ func (m *manager) SetRetryCount(count int) {
 }
 
 func (m *manager) AddConnectionObserver(observer ConnectionObserver) {
-	m.observers = append(m.observers, observer)
+	m.connectionObservers = append(m.connectionObservers, observer)
+}
+
+func (m *manager) setHeaderValues(_ *resty.Client, req *resty.Request) error {
+	req.SetHeaders(map[string]string{
+		"x-pm-appversion": m.cfg.AppVersion,
+		"User-Agent":      m.cfg.getUserAgent(),
+	})
+	return nil
 }
 
 func (m *manager) r(ctx context.Context) *resty.Request {
@@ -90,7 +112,7 @@ func (m *manager) handleRequestSuccess(_ *resty.Client, res *resty.Response) err
 
 	m.isDown = false
 
-	for _, observer := range m.observers {
+	for _, observer := range m.connectionObservers {
 		observer.OnUp()
 	}
 
@@ -113,15 +135,9 @@ func (m *manager) handleRequestFailure(req *resty.Request, err error) {
 
 	m.isDown = true
 
-	for _, observer := range m.observers {
+	for _, observer := range m.connectionObservers {
 		observer.OnDown()
 	}
 
 	go m.pingUntilSuccess()
 }
-
-func (m *manager) pingUntilSuccess() {
-	for m.testPing(context.Background()) != nil {
-		time.Sleep(time.Second) // TODO: How long to sleep here?
-	}
-}