From dd9651a94681ab9ee9774fd74c7fd66f3efe9afa Mon Sep 17 00:00:00 2001
From: Shautvast <shautvast@gmail.com>
Date: Fri, 20 Feb 2026 15:27:47 +0100
Subject: [PATCH] Add session expiry tracking and token refresh to
 proton-bridge
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Store expires_at (Unix timestamp) in session.json from ExpiresIn response field
- Add is_expired() with 5-minute refresh margin
- Implement POST /auth/v4/refresh flow: tries refresh before falling back to SRP login
- authenticate() now: use cached → refresh if expired → full login if all else fails

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 proton-bridge/src/auth.rs | 98 +++++++++++++++++++++++++++++++++++++--
 1 file changed, 93 insertions(+), 5 deletions(-)

diff --git a/proton-bridge/src/auth.rs b/proton-bridge/src/auth.rs
index c459986..a1c97bf 100644
--- a/proton-bridge/src/auth.rs
+++ b/proton-bridge/src/auth.rs
@@ -2,12 +2,15 @@ use reqwest::Client;
 use serde::{Deserialize, Serialize};
 use std::fs;
 use std::io::{self, Write};
+use std::time::{SystemTime, UNIX_EPOCH};
 
 use crate::config::ProtonConfig;
 use crate::srp;
 
 const API_BASE: &str = "https://mail.proton.me/api";
 const SESSION_FILE: &str = "session.json";
+/// Refresh 5 minutes before the token actually expires.
+const REFRESH_MARGIN_SECS: u64 = 300;
 
 // ── API types ───────────────────────────────────────────────────────────────
 
@@ -53,11 +56,37 @@ struct AuthResponse {
     refresh_token: String,
     #[serde(rename = "ServerProof")]
     server_proof: String,
+    #[serde(rename = "ExpiresIn")]
+    expires_in: u64,
     /// Integer bitmask: 0 = none, 1 = TOTP, 2 = FIDO2.
     #[serde(rename = "TwoFactor")]
     two_factor: u32,
 }
 
+#[derive(Serialize)]
+struct RefreshRequest<'a> {
+    #[serde(rename = "UID")]
+    uid: &'a str,
+    #[serde(rename = "RefreshToken")]
+    refresh_token: &'a str,
+    #[serde(rename = "ResponseType")]
+    response_type: &'a str,
+    #[serde(rename = "GrantType")]
+    grant_type: &'a str,
+    #[serde(rename = "RedirectURI")]
+    redirect_uri: &'a str,
+}
+
+#[derive(Deserialize)]
+struct RefreshResponse {
+    #[serde(rename = "AccessToken")]
+    access_token: String,
+    #[serde(rename = "RefreshToken")]
+    refresh_token: String,
+    #[serde(rename = "ExpiresIn")]
+    expires_in: u64,
+}
+
 #[derive(Serialize)]
 struct TotpRequest<'a> {
     #[serde(rename = "TwoFactorCode")]
@@ -71,9 +100,19 @@ pub struct Session {
     pub uid: String,
     pub access_token: String,
     pub refresh_token: String,
+    /// Unix timestamp at which the access token expires.
+    pub expires_at: u64,
 }
 
 impl Session {
+    pub fn is_expired(&self) -> bool {
+        let now = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap_or_default()
+            .as_secs();
+        now + REFRESH_MARGIN_SECS >= self.expires_at
+    }
+
     pub fn save(&self) -> Result<(), String> {
         let json = serde_json::to_string_pretty(self).map_err(|e| e.to_string())?;
         fs::write(SESSION_FILE, json).map_err(|e| e.to_string())
@@ -85,6 +124,14 @@ impl Session {
     }
 }
 
+fn expires_at_from_now(expires_in: u64) -> u64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap_or_default()
+        .as_secs()
+        + expires_in
+}
+
 // ── auth flow ────────────────────────────────────────────────────────────────
 
 /// Build a reqwest client with the headers Proton's API requires.
@@ -100,16 +147,56 @@ pub fn build_client() -> Result<Client, String> {
         .map_err(|e| e.to_string())
 }
 
-/// Authenticate against the Proton API and return a valid session.
-/// Loads a cached session from disk if available; falls back to full SRP login.
+/// Return a valid session, refreshing or re-logging in as needed.
 pub async fn authenticate(client: &Client, config: &ProtonConfig) -> Result<Session, String> {
     if let Some(session) = Session::load() {
-        println!("Loaded cached session for {}", config.username);
-        return Ok(session);
+        if !session.is_expired() {
+            println!("Using cached session for {}", config.username);
+            return Ok(session);
+        }
+        println!("Access token expired, refreshing…");
+        match refresh(client, &session).await {
+            Ok(refreshed) => {
+                refreshed.save()?;
+                println!("Session refreshed for {}", config.username);
+                return Ok(refreshed);
+            }
+            Err(e) => {
+                println!("Refresh failed ({}), falling back to full login", e);
+            }
+        }
     }
     login(client, config).await
 }
 
+async fn refresh(client: &Client, session: &Session) -> Result<Session, String> {
+    let resp: RefreshResponse = client
+        .post(format!("{}/auth/v4/refresh", API_BASE))
+        .header("x-pm-uid", &session.uid)
+        .json(&RefreshRequest {
+            uid: &session.uid,
+            refresh_token: &session.refresh_token,
+            response_type: "token",
+            grant_type: "refresh_token",
+            redirect_uri: "https://proton.me",
+        })
+        .send()
+        .await
+        .map_err(|e| e.to_string())?
+        .error_for_status()
+        .map_err(|e| e.to_string())?
+        .json()
+        .await
+        .map_err(|e| e.to_string())?;
+
+    Ok(Session {
+        uid: session.uid.clone(),
+        access_token: resp.access_token,
+        refresh_token: resp.refresh_token,
+        expires_at: expires_at_from_now(resp.expires_in),
+    })
+}
+
 async fn login(client: &Client, config: &ProtonConfig) -> Result<Session, String> {
     // Step 1: request SRP challenge
     let info: AuthInfoResponse = client
@@ -162,13 +249,14 @@ async fn login(client: &Client, config: &ProtonConfig) -> Result<Session, String
         return Err("Server proof verification failed — possible MITM".to_string());
     }
 
-    // Step 4: handle TOTP if enabled
     let session = Session {
         uid: auth.uid.clone(),
         access_token: auth.access_token.clone(),
         refresh_token: auth.refresh_token.clone(),
+        expires_at: expires_at_from_now(auth.expires_in),
     };
 
+    // Step 4: handle TOTP if enabled
     if auth.two_factor & 1 != 0 {
         submit_totp(client, &session).await?;
     }