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 serde::{Deserialize, Serialize};
|
||||||
use std::fs;
|
use std::fs;
|
||||||
use std::io::{self, Write};
|
use std::io::{self, Write};
|
||||||
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
use crate::config::ProtonConfig;
|
use crate::config::ProtonConfig;
|
||||||
use crate::srp;
|
use crate::srp;
|
||||||
|
|
||||||
const API_BASE: &str = "https://mail.proton.me/api";
|
const API_BASE: &str = "https://mail.proton.me/api";
|
||||||
const SESSION_FILE: &str = "session.json";
|
const SESSION_FILE: &str = "session.json";
|
||||||
|
/// Refresh 5 minutes before the token actually expires.
|
||||||
|
const REFRESH_MARGIN_SECS: u64 = 300;
|
||||||
|
|
||||||
// ── API types ───────────────────────────────────────────────────────────────
|
// ── API types ───────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
@ -53,11 +56,37 @@ struct AuthResponse {
|
||||||
refresh_token: String,
|
refresh_token: String,
|
||||||
#[serde(rename = "ServerProof")]
|
#[serde(rename = "ServerProof")]
|
||||||
server_proof: String,
|
server_proof: String,
|
||||||
|
#[serde(rename = "ExpiresIn")]
|
||||||
|
expires_in: u64,
|
||||||
/// Integer bitmask: 0 = none, 1 = TOTP, 2 = FIDO2.
|
/// Integer bitmask: 0 = none, 1 = TOTP, 2 = FIDO2.
|
||||||
#[serde(rename = "TwoFactor")]
|
#[serde(rename = "TwoFactor")]
|
||||||
two_factor: u32,
|
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)]
|
#[derive(Serialize)]
|
||||||
struct TotpRequest<'a> {
|
struct TotpRequest<'a> {
|
||||||
#[serde(rename = "TwoFactorCode")]
|
#[serde(rename = "TwoFactorCode")]
|
||||||
|
|
@ -71,9 +100,19 @@ pub struct Session {
|
||||||
pub uid: String,
|
pub uid: String,
|
||||||
pub access_token: String,
|
pub access_token: String,
|
||||||
pub refresh_token: String,
|
pub refresh_token: String,
|
||||||
|
/// Unix timestamp at which the access token expires.
|
||||||
|
pub expires_at: u64,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Session {
|
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> {
|
pub fn save(&self) -> Result<(), String> {
|
||||||
let json = serde_json::to_string_pretty(self).map_err(|e| e.to_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())
|
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 ────────────────────────────────────────────────────────────────
|
// ── auth flow ────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
/// Build a reqwest client with the headers Proton's API requires.
|
/// 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())
|
.map_err(|e| e.to_string())
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Authenticate against the Proton API and return a valid session.
|
/// Return a valid session, refreshing or re-logging in as needed.
|
||||||
/// Loads a cached session from disk if available; falls back to full SRP login.
|
|
||||||
pub async fn authenticate(client: &Client, config: &ProtonConfig) -> Result<Session, String> {
|
pub async fn authenticate(client: &Client, config: &ProtonConfig) -> Result<Session, String> {
|
||||||
if let Some(session) = Session::load() {
|
if let Some(session) = Session::load() {
|
||||||
println!("Loaded cached session for {}", config.username);
|
if !session.is_expired() {
|
||||||
return Ok(session);
|
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
|
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> {
|
async fn login(client: &Client, config: &ProtonConfig) -> Result<Session, String> {
|
||||||
// Step 1: request SRP challenge
|
// Step 1: request SRP challenge
|
||||||
let info: AuthInfoResponse = client
|
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());
|
return Err("Server proof verification failed — possible MITM".to_string());
|
||||||
}
|
}
|
||||||
|
|
||||||
// Step 4: handle TOTP if enabled
|
|
||||||
let session = Session {
|
let session = Session {
|
||||||
uid: auth.uid.clone(),
|
uid: auth.uid.clone(),
|
||||||
access_token: auth.access_token.clone(),
|
access_token: auth.access_token.clone(),
|
||||||
refresh_token: auth.refresh_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 {
|
if auth.two_factor & 1 != 0 {
|
||||||
submit_totp(client, &session).await?;
|
submit_totp(client, &session).await?;
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue