Refresh token logic

2 files changed, 76 insertions, 39 deletions

diff --git a/backend/auth/oauth.go b/backend/auth/oauth.go
index b217ad2..9be1da0 100644
--- a/backend/auth/oauth.go
+++ b/backend/auth/oauth.go
@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
+	"time"
 )
 
 type OAuthConfig struct {
@@ -25,6 +26,7 @@ type OAuthState struct {
 	PDS          string
 	AccessToken  string
 	RefreshToken string
+	ExpiresAt    time.Time
 	KeyPair      *DPoPKeyPair
 }
 
@@ -110,32 +112,16 @@ func BuildAuthorizationURL(metadata *OAuthMetadata, config *OAuthConfig, state,
 	return metadata.AuthorizationEndpoint + "?" + params.Encode()
 }
 
-func ExchangeCodeForTokens(pdsURL, code, verifier, backendURL string, keypair *DPoPKeyPair) (*TokenResponse, error) {
-	metadata, err := FetchOAuthMetadata(pdsURL)
-	if err != nil {
-		return nil, err
-	}
-
-	data := url.Values{
-		"grant_type":    {"authorization_code"},
-		"code":          {code},
-		"redirect_uri":  {backendURL + "/api/callback"},
-		"code_verifier": {verifier},
-		"client_id":     {backendURL + "/api/callback"},
-	}
-
-	fmt.Printf("Sending verifier: %s\n", verifier)
-
-	dpopProof, err := keypair.GenerateDPoPProof("POST", metadata.TokenEndpoint, "")
+func postWithDPoP(url string, data url.Values, keypair *DPoPKeyPair) (*TokenResponse, error) {
+	dpopProof, err := keypair.GenerateDPoPProof("POST", url, "")
 	if err != nil {
 		return nil, err
 	}
 
-	req, err := http.NewRequest("POST", metadata.TokenEndpoint, strings.NewReader(data.Encode()))
+	req, err := http.NewRequest("POST", url, strings.NewReader(data.Encode()))
 	if err != nil {
 		return nil, err
 	}
-
 	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
 	req.Header.Set("DPoP", dpopProof)
 
@@ -145,7 +131,6 @@ func ExchangeCodeForTokens(pdsURL, code, verifier, backendURL string, keypair *D
 	}
 	defer resp.Body.Close()
 
-	// if nonce required, retry with nonce
 	if resp.StatusCode == 400 {
 		body, _ := io.ReadAll(resp.Body)
 		bodyStr := string(body)
@@ -156,13 +141,12 @@ func ExchangeCodeForTokens(pdsURL, code, verifier, backendURL string, keypair *D
 				return nil, fmt.Errorf("server requires nonce but didn't provide one")
 			}
 
-			// generate new DPoP proof with nonce
-			dpopProof, err = keypair.GenerateDPoPProof("POST", metadata.TokenEndpoint, nonce)
+			dpopProof, err = keypair.GenerateDPoPProof("POST", url, nonce)
 			if err != nil {
 				return nil, err
 			}
 
-			retryReq, err := http.NewRequest("POST", metadata.TokenEndpoint, strings.NewReader(data.Encode()))
+			retryReq, err := http.NewRequest("POST", url, strings.NewReader(data.Encode()))
 			if err != nil {
 				return nil, err
 			}
@@ -174,16 +158,14 @@ func ExchangeCodeForTokens(pdsURL, code, verifier, backendURL string, keypair *D
 				return nil, err
 			}
 			defer resp.Body.Close()
-		}  else {
-			// it's a different 400 error, don't retry
-			return nil, fmt.Errorf("token exchange failed: 400 - %s", bodyStr)
+		} else {
+			return nil, fmt.Errorf("request failed: 400 - %s", bodyStr)
 		}
 	}
 
 	if resp.StatusCode != 200 {
 		body, _ := io.ReadAll(resp.Body)
-		fmt.Printf("Token exchange error: %s\n", string(body))
-		return nil, fmt.Errorf("token exchange failed: %d - %s", resp.StatusCode, string(body))
+		return nil, fmt.Errorf("request failed: %d - %s", resp.StatusCode, string(body))
 	}
 
 	var tokens TokenResponse
@@ -193,3 +175,35 @@ func ExchangeCodeForTokens(pdsURL, code, verifier, backendURL string, keypair *D
 
 	return &tokens, nil
 }
+
+func ExchangeCodeForTokens(pdsURL, code, verifier, callbackURL string, keypair *DPoPKeyPair) (*TokenResponse, error) {
+	metadata, err := FetchOAuthMetadata(pdsURL)
+	if err != nil {
+		return nil, err
+	}
+
+	data := url.Values{
+		"grant_type":    {"authorization_code"},
+		"code":          {code},
+		"redirect_uri":  {callbackURL},
+		"code_verifier": {verifier},
+		"client_id":     {callbackURL},
+	}
+
+	return postWithDPoP(metadata.TokenEndpoint, data, keypair)
+}
+
+func RefreshAccessToken(pdsURL, refreshToken, callbackURL string, keypair *DPoPKeyPair) (*TokenResponse, error) {
+	metadata, err := FetchOAuthMetadata(pdsURL)
+	if err != nil {
+		return nil, err
+	}
+
+	data := url.Values{
+		"grant_type":    {"refresh_token"},
+		"refresh_token": {refreshToken},
+		"client_id":     {callbackURL},
+	}
+
+	return postWithDPoP(metadata.TokenEndpoint, data, keypair)
+}
diff --git a/backend/main.go b/backend/main.go
index 6598266..032c6e9 100644
--- a/backend/main.go
+++ b/backend/main.go
@@ -9,6 +9,7 @@ import (
 	"os"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/joho/godotenv"
 )
@@ -175,6 +176,7 @@ func callbackHandler(w http.ResponseWriter, r *http.Request) {
 		PDS:          oauthState.PDS,
 		AccessToken:  tokens.AccessToken,
 		RefreshToken: tokens.RefreshToken,
+		ExpiresAt:    time.Now().Add(time.Duration(tokens.ExpiresIn) * time.Second),
 		KeyPair:      oauthState.KeyPair,
 	}
 	delete(sessions, state) // remove the partial oauth state from login
@@ -243,21 +245,42 @@ func postsHandler(w http.ResponseWriter, r *http.Request) {
 
 	sessionsMutex.RLock()
 	session, exists := sessions[cookie.Value]
+	if !exists {
+		sessionsMutex.RUnlock()
+		http.Error(w, "invalid session", http.StatusUnauthorized)
+		return
+	}
+
 	fmt.Printf("Session exists: %v\n", exists)
-	if exists {
-		fmt.Printf("Session DID: %s\n", session.DID)
-		fmt.Printf("Session PDS: %s\n", session.PDS)
-		fmt.Printf("Session AccessToken length: %d\n", len(session.AccessToken))
-		fmt.Printf("Session KeyPair: %p\n", session.KeyPair)
-		if session.KeyPair != nil {
-			fmt.Printf("KeyPair PublicKey: %p\n", session.KeyPair.PublicKey)
-		}
+	fmt.Printf("Session DID: %s\n", session.DID)
+	fmt.Printf("Session PDS: %s\n", session.PDS)
+	fmt.Printf("Session AccessToken length: %d\n", len(session.AccessToken))
+	fmt.Printf("Session KeyPair: %p\n", session.KeyPair)
+	if session.KeyPair != nil {
+		fmt.Printf("KeyPair PublicKey: %p\n", session.KeyPair.PublicKey)
 	}
+
+	needsRefresh := time.Now().After(session.ExpiresAt)
+	pds := session.PDS
+	refreshToken := session.RefreshToken
+	keypair := session.KeyPair
 	sessionsMutex.RUnlock()
 
-	if !exists {
-		http.Error(w, "invalid session", http.StatusUnauthorized)
-		return
+	if needsRefresh {
+		newTokens, err := auth.RefreshAccessToken(pds, refreshToken, keypair)
+		if err != nil {
+			sessionsMutex.Lock()
+			delete(sessions, cookie.Value)
+			sessionsMutex.Unlock()
+			http.Error(w, "session expired, please login again", http.StatusUnauthorized)
+			return
+		}
+
+		sessionsMutex.Lock()
+		session.AccessToken = newTokens.AccessToken
+		session.RefreshToken = newTokens.RefreshToken
+		session.ExpiresAt = time.Now().Add(time.Duration(newTokens.ExpiresIn) * time.Second)
+		sessionsMutex.Unlock()
 	}
 
 	posts, err := fetchers.FetchAppBskyFeedPosts(session.DID, session.PDS, session.AccessToken, session.KeyPair)