From cc03e059b9df825152899f0f17d8bba2e209fb40 Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Sat, 31 Aug 2019 17:43:46 +0300
Subject: [PATCH] Save key salts after auth

The key salts endpoint can only be accessed right after authentication, it's
blocked after auth refresh.

Closes: https://github.com/emersion/hydroxide/issues/68
---
 auth/auth.go               |  3 ++-
 cmd/hydroxide/hydroxide.go | 14 ++++++++++----
 protonmail/auth.go         | 16 +++++++---------
 3 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/auth/auth.go b/auth/auth.go
index 6d5c770..bf6c568 100644
--- a/auth/auth.go
+++ b/auth/auth.go
@@ -25,6 +25,7 @@ type CachedAuth struct {
 	protonmail.Auth
 	LoginPassword   string
 	MailboxPassword string
+	KeySalts        map[string][]byte
 	// TODO: add padding
 }
 
@@ -131,7 +132,7 @@ func authenticate(c *protonmail.Client, cachedAuth *CachedAuth, username string)
 	}
 	cachedAuth.Auth = *auth
 
-	return c.Unlock(auth, cachedAuth.MailboxPassword)
+	return c.Unlock(auth, cachedAuth.KeySalts, cachedAuth.MailboxPassword)
 }
 
 func ListUsernames() ([]string, error) {
diff --git a/cmd/hydroxide/hydroxide.go b/cmd/hydroxide/hydroxide.go
index fa601c2..a7c5b9e 100644
--- a/cmd/hydroxide/hydroxide.go
+++ b/cmd/hydroxide/hydroxide.go
@@ -184,7 +184,12 @@ func main() {
 			}
 		}
 
-		_, err := c.Unlock(a, mailboxPassword)
+		keySalts, err := c.ListKeySalts()
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		_, err = c.Unlock(a, keySalts, mailboxPassword)
 		if err != nil {
 			log.Fatal(err)
 		}
@@ -195,9 +200,10 @@ func main() {
 		}
 
 		err = auth.EncryptAndSave(&auth.CachedAuth{
-			*a,
-			loginPassword,
-			mailboxPassword,
+			Auth: *a,
+			LoginPassword: loginPassword,
+			MailboxPassword: mailboxPassword,
+			KeySalts: keySalts,
 		}, username, secretKey)
 		if err != nil {
 			log.Fatal(err)
diff --git a/protonmail/auth.go b/protonmail/auth.go
index dbf1eaa..dd09fe3 100644
--- a/protonmail/auth.go
+++ b/protonmail/auth.go
@@ -141,7 +141,10 @@ func (c *Client) Auth(username, password, twoFactorCode string, info *AuthInfo)
 		return nil, err
 	}
 
-	return respData.auth(), nil
+	auth := respData.auth()
+	c.uid = auth.UID
+	c.accessToken = auth.AccessToken
+	return auth, nil
 }
 
 type authRefreshReq struct {
@@ -178,7 +181,7 @@ func (c *Client) AuthRefresh(expiredAuth *Auth) (*Auth, error) {
 	return auth, nil
 }
 
-func (c *Client) listKeySalts() (map[string][]byte, error) {
+func (c *Client) ListKeySalts() (map[string][]byte, error) {
 	req, err := c.newRequest(http.MethodGet, "/keys/salts", nil)
 	if err != nil {
 		return nil, err
@@ -235,7 +238,7 @@ func unlockKey(e *openpgp.Entity, passphraseBytes []byte) error {
 	return nil
 }
 
-func (c *Client) Unlock(auth *Auth, passphrase string) (openpgp.EntityList, error) {
+func (c *Client) Unlock(auth *Auth, keySalts map[string][]byte, passphrase string) (openpgp.EntityList, error) {
 	c.uid = auth.UID
 	c.accessToken = auth.AccessToken
 
@@ -244,11 +247,6 @@ func (c *Client) Unlock(auth *Auth, passphrase string) (openpgp.EntityList, erro
 		return nil, err
 	}
 
-	salts, err := c.listKeySalts()
-	if err != nil {
-		return nil, err
-	}
-
 	var keyRing openpgp.EntityList
 	for _, addr := range addrs {
 		for _, key := range addr.Keys {
@@ -259,7 +257,7 @@ func (c *Client) Unlock(auth *Auth, passphrase string) (openpgp.EntityList, erro
 			}
 
 			passphraseBytes := []byte(passphrase)
-			if keySalt, ok := salts[key.ID]; ok && keySalt != nil {
+			if keySalt, ok := keySalts[key.ID]; ok && keySalt != nil {
 				passphraseBytes, err = computeKeyPassword(passphraseBytes, keySalt)
 				if err != nil {
 					return nil, err