From 47858885903ff92c9a7ee3715296ff26f1e57be0 Mon Sep 17 00:00:00 2001
From: emersion <contact@emersion.fr>
Date: Tue, 22 Aug 2017 12:23:58 +0200
Subject: [PATCH] Change Client.AuthRefresh to take an Auth

---
 cmd/hydroxide/hydroxide.go |  6 +-----
 protonmail/auth.go         | 16 +++++++++++-----
 protonmail/protonmail.go   |  5 +++++
 3 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/cmd/hydroxide/hydroxide.go b/cmd/hydroxide/hydroxide.go
index bf847ce..9b7e553 100644
--- a/cmd/hydroxide/hydroxide.go
+++ b/cmd/hydroxide/hydroxide.go
@@ -47,15 +47,11 @@ func main() {
 	var password string
 	auth, err := readCachedAuth()
 	if err == nil {
-		passwordMode := auth.PasswordMode
-
 		var err error
-		auth, err = c.AuthRefresh(auth.UID, auth.RefreshToken)
+		auth, err = c.AuthRefresh(auth)
 		if err != nil {
 			log.Fatal(err)
 		}
-
-		auth.PasswordMode = passwordMode
 	} else if os.IsNotExist(err) {
 		fmt.Printf("Username: ")
 		scanner.Scan()
diff --git a/protonmail/auth.go b/protonmail/auth.go
index 24e7354..3b27054 100644
--- a/protonmail/auth.go
+++ b/protonmail/auth.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"net/http"
 	"strings"
+	"time"
 
 	"golang.org/x/crypto/openpgp"
 )
@@ -82,7 +83,7 @@ const (
 )
 
 type Auth struct {
-	ExpiresIn    int
+	ExpiresAt    time.Time
 	Scope        string
 	UID          string `json:"Uid"`
 	RefreshToken string
@@ -97,6 +98,7 @@ type Auth struct {
 type authResp struct {
 	resp
 	Auth
+	ExpiresIn   int
 	AccessToken string
 	TokenType   string
 	ServerProof string
@@ -106,6 +108,7 @@ type authResp struct {
 
 func (resp *authResp) auth() *Auth {
 	auth := &resp.Auth
+	auth.ExpiresAt = time.Now().Add(time.Duration(resp.ExpiresIn) * time.Second)
 	auth.accessToken = resp.AccessToken
 	auth.privateKey = resp.PrivateKey
 	auth.keySalt = resp.KeySalt
@@ -164,11 +167,11 @@ type authRefreshReq struct {
 	State string
 }
 
-func (c *Client) AuthRefresh(uid, refreshToken string) (*Auth, error) {
+func (c *Client) AuthRefresh(expiredAuth *Auth) (*Auth, error) {
 	reqData := &authRefreshReq{
 		ClientID:        c.ClientID,
-		UID:             uid,
-		RefreshToken:    refreshToken,
+		UID:             expiredAuth.UID,
+		RefreshToken:    expiredAuth.RefreshToken,
 	}
 
 	req, err := c.newJSONRequest(http.MethodPost, "/auth/refresh", reqData)
@@ -181,7 +184,10 @@ func (c *Client) AuthRefresh(uid, refreshToken string) (*Auth, error) {
 		return nil, err
 	}
 
-	return respData.auth(), nil
+	auth := respData.auth()
+	//auth.EventID = expiredAuth.EventID
+	auth.PasswordMode = expiredAuth.PasswordMode
+	return auth, nil
 }
 
 func (c *Client) Unlock(auth *Auth, passphrase string) (openpgp.EntityList, error) {
diff --git a/protonmail/protonmail.go b/protonmail/protonmail.go
index d797e0f..10df4c9 100644
--- a/protonmail/protonmail.go
+++ b/protonmail/protonmail.go
@@ -63,6 +63,11 @@ func (c *Client) newRequest(method, path string, body io.Reader) (*http.Request,
 	req.Header.Set("X-Pm-Appversion", c.AppVersion)
 	req.Header.Set(headerAPIVersion, strconv.Itoa(Version))
 
+	if c.uid != "" && c.accessToken != "" {
+		req.Header.Set("X-Pm-Uid", c.uid)
+		req.Header.Set("Authorization", "Bearer " + c.accessToken)
+	}
+
 	return req, nil
 }