feat: persistent cookies

Changelog.md

@@ -4,6 +4,9 @@ Changelog [format](http://keepachangelog.com/en/1.0.0/)
 
 ## Unreleased
 
+### Added
+* GODT-633 Persistent anonymous API cookies for better load balancing and abuse detection. 
+
 ### Changed
 * GODT-409 Set flags have to replace all flags.
 * GODT-531 Better way to add trusted certificate in macOS.

cmd/Desktop-Bridge/main.go

@@ -56,6 +56,7 @@ import (
 	"github.com/ProtonMail/proton-bridge/pkg/args"
 	"github.com/ProtonMail/proton-bridge/pkg/config"
 	"github.com/ProtonMail/proton-bridge/pkg/constants"
+	"github.com/ProtonMail/proton-bridge/pkg/cookies"
 	"github.com/ProtonMail/proton-bridge/pkg/listener"
 	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
 	"github.com/ProtonMail/proton-bridge/pkg/updates"
@@ -273,6 +274,13 @@ func run(context *cli.Context) (contextError error) { // nolint[funlen]
 	// implementation depending on whether build flag pmapi_prod is used or not.
 	cm.SetRoundTripper(cfg.GetRoundTripper(cm, eventListener))
 
+	// Cookies must be persisted across restarts.
+	jar, err := cookies.New(cookies.NewPersister(pref))
+	if err != nil {
+		logrus.WithError(err).Fatal("Could not create cookie jar")
+	}
+	cm.SetCookieJar(jar)
+
 	bridgeInstance := bridge.New(cfg, pref, panicHandler, eventListener, cm, credentialsStore)
 	imapBackend := imap.NewIMAPBackend(panicHandler, eventListener, cfg, bridgeInstance)
 	smtpBackend := smtp.NewSMTPBackend(panicHandler, eventListener, pref, bridgeInstance)

internal/bridge/release_notes.go

@@ -15,7 +15,7 @@
 // You should have received a copy of the GNU General Public License
 // along with ProtonMail Bridge.  If not, see <https://www.gnu.org/licenses/>.
 
-// Code generated by ./release-notes.sh at Thu 21 May 2020 07:59:59 AM CEST. DO NOT EDIT.
+// Code generated by ./release-notes.sh at Wed 12 Aug 2020 12:25:26 PM CEST. DO NOT EDIT.
 
 package bridge
 

internal/preferences/preferences.go

@@ -37,6 +37,7 @@ const (
 	SMTPSSLKey             = "user_ssl_smtp"
 	AllowProxyKey          = "allow_proxy"
 	AutostartKey           = "autostart"
+	CookiesKey             = "cookies"
 	ReportOutgoingNoEncKey = "report_outgoing_email_without_encryption"
 	LastVersionKey         = "last_used_version"
 )

pkg/cookies/jar.go

@@ -0,0 +1,61 @@
+package cookies
+
+import (
+	"net/http"
+	"net/http/cookiejar"
+	"net/url"
+	"sync"
+
+	"github.com/sirupsen/logrus"
+)
+
+type Jar struct {
+	jar       *cookiejar.Jar
+	persister *Persister
+	locker    sync.Locker
+}
+
+func New(persister *Persister) (*Jar, error) {
+	jar, err := cookiejar.New(nil)
+	if err != nil {
+		return nil, err
+	}
+
+	cookies, err := persister.Load()
+	if err != nil {
+		return nil, err
+	}
+
+	for rawURL, cookies := range cookies {
+		url, err := url.Parse(rawURL)
+		if err != nil {
+			continue
+		}
+
+		jar.SetCookies(url, cookies)
+	}
+
+	return &Jar{
+		jar:       jar,
+		persister: persister,
+		locker:    &sync.Mutex{},
+	}, nil
+}
+
+func (j *Jar) SetCookies(u *url.URL, cookies []*http.Cookie) {
+	j.locker.Lock()
+	defer j.locker.Unlock()
+
+	j.jar.SetCookies(u, cookies)
+
+	if err := j.persister.Persist(u.String(), cookies); err != nil {
+		logrus.WithError(err).Warn("Failed to persist cookie")
+	}
+}
+
+func (j *Jar) Cookies(u *url.URL) []*http.Cookie {
+	j.locker.Lock()
+	defer j.locker.Unlock()
+
+	return j.jar.Cookies(u)
+}

pkg/cookies/jar_test.go

@@ -0,0 +1,80 @@
+package cookies
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestJar(t *testing.T) {
+	testCookies := []testCookie{
+		{"TestName1", "TestValue1"},
+		{"TestName2", "TestValue2"},
+		{"TestName3", "TestValue3"},
+	}
+
+	ts := getTestServer(t, testCookies...)
+	defer ts.Close()
+
+	jar, err := New(NewPersister(make(testPersister)))
+	require.NoError(t, err)
+
+	client := &http.Client{Jar: jar}
+
+	setRes, err := client.Get(ts.URL + "/set")
+	if err != nil {
+		t.FailNow()
+	}
+	require.NoError(t, setRes.Body.Close())
+
+	getRes, err := client.Get(ts.URL + "/get")
+	if err != nil {
+		t.FailNow()
+	}
+	require.NoError(t, getRes.Body.Close())
+}
+
+type testCookie struct {
+	name, value string
+}
+
+func getTestServer(t *testing.T, wantCookies ...testCookie) *httptest.Server {
+	mux := http.NewServeMux()
+
+	mux.HandleFunc("/set", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		for _, cookie := range wantCookies {
+			http.SetCookie(w, &http.Cookie{
+				Name:  cookie.name,
+				Value: cookie.value,
+			})
+		}
+
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	mux.HandleFunc("/get", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		require.Len(t, r.Cookies(), len(wantCookies))
+
+		for k, v := range r.Cookies() {
+			assert.Equal(t, wantCookies[k].name, v.Name)
+			assert.Equal(t, wantCookies[k].value, v.Value)
+		}
+
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	return httptest.NewServer(mux)
+}
+
+type testPersister map[string]string
+
+func (p testPersister) Set(key, value string) {
+	p[key] = value
+}
+
+func (p testPersister) Get(key string) string {
+	return p[key]
+}

pkg/cookies/persister.go

@@ -0,0 +1,91 @@
+package cookies
+
+import (
+	"encoding/json"
+	"net/http"
+
+	"github.com/ProtonMail/proton-bridge/internal/preferences"
+)
+
+type Persister struct {
+	prefs GetterSetter
+}
+
+type GetterSetter interface {
+	Get(string) string
+	Set(string, string)
+}
+
+func NewPersister(prefs GetterSetter) *Persister {
+	return &Persister{prefs: prefs}
+}
+
+func (p *Persister) Persist(url string, cookies []*http.Cookie) error {
+	b, err := json.Marshal(cookies)
+	if err != nil {
+		return err
+	}
+
+	val, err := p.load()
+	if err != nil {
+		return err
+	}
+
+	val[url] = string(b)
+
+	return p.save(val)
+}
+
+func (p *Persister) Load() (map[string][]*http.Cookie, error) {
+	res := make(map[string][]*http.Cookie)
+
+	val, err := p.load()
+	if err != nil {
+		return nil, err
+	}
+
+	for url, rawCookies := range val {
+		var cookies []*http.Cookie
+
+		if err := json.Unmarshal([]byte(rawCookies), &cookies); err != nil {
+			return nil, err
+		}
+
+		res[url] = cookies
+	}
+
+	return res, nil
+}
+
+type dataStructure map[string]string
+
+func (p *Persister) load() (dataStructure, error) {
+	b := p.prefs.Get(preferences.CookiesKey)
+
+	if b == "" {
+		if err := p.save(make(dataStructure)); err != nil {
+			return nil, err
+		}
+
+		return p.load()
+	}
+
+	var val dataStructure
+
+	if err := json.Unmarshal([]byte(b), &val); err != nil {
+		return nil, err
+	}
+
+	return val, nil
+}
+
+func (p *Persister) save(val dataStructure) error {
+	b, err := json.Marshal(val)
+	if err != nil {
+		return err
+	}
+
+	p.prefs.Set(preferences.CookiesKey, string(b))
+
+	return nil
+}

pkg/pmapi/client.go

@@ -127,7 +127,7 @@ type client struct {
 func newClient(cm *ClientManager, userID string) *client {
 	return &client{
 		cm:            cm,
-		hc:            getHTTPClient(cm.config, cm.roundTripper),
+		hc:            getHTTPClient(cm.config, cm.roundTripper, cm.cookieJar),
 		userID:        userID,
 		requestLocker: &sync.Mutex{},
 		keyRingLock:   &sync.Mutex{},
@@ -137,10 +137,11 @@ func newClient(cm *ClientManager, userID string) *client {
 }
 
 // getHTTPClient returns a http client configured by the given client config and using the given transport.
-func getHTTPClient(cfg *ClientConfig, rt http.RoundTripper) (hc *http.Client) {
+func getHTTPClient(cfg *ClientConfig, rt http.RoundTripper, jar http.CookieJar) (hc *http.Client) {
 	return &http.Client{
-		Timeout:   cfg.Timeout,
 		Transport: rt,
+		Jar:       jar,
+		Timeout:   cfg.Timeout,
 	}
 }
 

pkg/pmapi/clientmanager.go

@@ -41,6 +41,7 @@ type ClientManager struct {
 
 	clients       map[string]Client
 	clientsLocker sync.Locker
+	cookieJar     http.CookieJar
 
 	tokens       map[string]string
 	tokensLocker sync.Locker
@@ -126,6 +127,11 @@ func (cm *ClientManager) SetClientConstructor(f func(userID string) Client) {
 	cm.newClient = f
 }
 
+// SetCookieJar sets the cookie jar given to clients.
+func (cm *ClientManager) SetCookieJar(jar http.CookieJar) {
+	cm.cookieJar = jar
+}
+
 // SetRoundTripper sets the roundtripper used by clients created by this client manager.
 func (cm *ClientManager) SetRoundTripper(rt http.RoundTripper) {
 	cm.roundTripper = rt
@@ -145,9 +151,11 @@ func (cm *ClientManager) GetClient(userID string) Client {
 		return client
 	}
 
-	cm.clients[userID] = cm.newClient(userID)
+	client := cm.newClient(userID)
 
-	return cm.clients[userID]
+	cm.clients[userID] = client
+
+	return client
 }
 
 // GetAnonymousClient returns an anonymous client.
@@ -303,7 +311,7 @@ var ErrNoInternetConnection = errors.New("no internet connection")
 // CheckConnection returns an error if there is no internet connection.
 // This should be moved to the ConnectionManager when it is implemented.
 func (cm *ClientManager) CheckConnection() error {
-	client := getHTTPClient(cm.config, cm.roundTripper)
+	client := getHTTPClient(cm.config, cm.roundTripper, cm.cookieJar)
 
 	// Do not cumulate timeouts, use goroutines.
 	retStatus := make(chan error)