diff --git a/Changelog.md b/Changelog.md index 287a5eb6..f998c6af 100644 --- a/Changelog.md +++ b/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. diff --git a/cmd/Desktop-Bridge/main.go b/cmd/Desktop-Bridge/main.go index fafa6baa..1099c6ba 100644 --- a/cmd/Desktop-Bridge/main.go +++ b/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) diff --git a/internal/bridge/release_notes.go b/internal/bridge/release_notes.go index 02a66d94..a5f4e263 100644 --- a/internal/bridge/release_notes.go +++ b/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 . -// 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 diff --git a/internal/preferences/preferences.go b/internal/preferences/preferences.go index e1767c49..cba4f864 100644 --- a/internal/preferences/preferences.go +++ b/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" ) diff --git a/pkg/cookies/jar.go b/pkg/cookies/jar.go new file mode 100644 index 00000000..afe9a57f --- /dev/null +++ b/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) +} diff --git a/pkg/cookies/jar_test.go b/pkg/cookies/jar_test.go new file mode 100644 index 00000000..472353b1 --- /dev/null +++ b/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] +} diff --git a/pkg/cookies/persister.go b/pkg/cookies/persister.go new file mode 100644 index 00000000..b911c59c --- /dev/null +++ b/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 +} diff --git a/pkg/pmapi/client.go b/pkg/pmapi/client.go index b483deb8..8c76658e 100644 --- a/pkg/pmapi/client.go +++ b/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, } } diff --git a/pkg/pmapi/clientmanager.go b/pkg/pmapi/clientmanager.go index 6159f4f8..c4e44d80 100644 --- a/pkg/pmapi/clientmanager.go +++ b/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)