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:
parent
05b6aac692
commit
dd9651a946
1 changed files with 93 additions and 5 deletions
|
|
@ -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);
|
||||
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?;
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue