Add session expiry tracking and token refresh to proton-bridge

- 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>
This commit is contained in:
Shautvast 2026-02-20 15:27:47 +01:00
parent 05b6aac692
commit dd9651a946

View file

@ -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?;
}