fix(stockbit): keep a working session through a failed login check

- login restored the previous session on any check error, even after the
  check had rotated the pasted session: a timeout then left no working
  session. It now keeps a rotated session unless Stockbit refused it.
- The browser sign-in marked a cookie refused on any error, so a 429 or a
  timeout held a signed-in user until the wait ran out. Passing errors
  are retried after a pause; only a refusal is final.
- The local page checked concurrent pastes on their own threads, so a
  later one could overwrite the session the first had just rotated.
  Pastes are now checked one at a time, and none after one is accepted.

Claude-Session: https://claude.ai/code/session_01RqPPVeVDrnLoP7NUZmLRxY
This commit is contained in:
Rasyidan Akbar F. 2026-11-04 17:31:20 +07:00
commit e7f1fa2c90
3 changed files with 139 additions and 6 deletions

View file

@ -55,6 +55,12 @@ pub const PROFILE_DIR: &str = "stockbit-browser";
const START_TIMEOUT: Duration = Duration::from_secs(20);
/// Pause between cookie reads.
const POLL: Duration = Duration::from_millis(500);
/// How long a session check that failed for a passing reason (network, rate
/// limit) waits before the same cookie is checked again.
#[cfg(not(test))]
const VERIFY_RETRY: Duration = Duration::from_secs(5);
#[cfg(test)]
const VERIFY_RETRY: Duration = Duration::from_millis(50);
/// Socket timeout for one DevTools message.
const IO_TIMEOUT: Duration = Duration::from_secs(10);
/// How long the browser gets to quit after `Browser.close`.
@ -228,6 +234,8 @@ fn wait_for_session(
) -> Result<(), IdxError> {
let deadline = Instant::now() + timeout;
let mut rejected: Option<String> = None;
// After a check fails for a passing reason, the next is held back.
let mut retry_at: Option<Instant> = None;
let mut last_error: Option<String> = None;
// The window is watched so closing it ends the wait instead of running
// out the clock; the first polls may not see a page yet.
@ -235,7 +243,10 @@ fn wait_for_session(
let mut misses = 0;
loop {
match devtools.session_cookie() {
Ok(Some(cookie)) if rejected.as_deref() != Some(cookie.value.as_str()) => {
Ok(Some(cookie))
if rejected.as_deref() != Some(cookie.value.as_str())
&& retry_at.is_none_or(|at| Instant::now() >= at) =>
{
match verify(&cookie.value) {
Ok(()) => {
// The session now lives in idx's own file. Take the
@ -250,12 +261,20 @@ fn wait_for_session(
devtools.close();
return Ok(());
}
Err(err) if super::check_is_transient(&err) => {
crate::runtime::warn(format!(
"could not check the session in idx's browser yet ({err}); \
retrying"
));
retry_at = Some(Instant::now() + VERIFY_RETRY);
}
Err(err) => {
crate::runtime::warn(format!(
"Stockbit refused the session in idx's browser ({err}); sign in \
again in that window"
));
rejected = Some(cookie.value);
retry_at = None;
}
}
}
@ -979,6 +998,33 @@ mod tests {
let _ = std::fs::remove_dir_all(&profile);
}
/// A check that fails for a passing reason (a timeout, a 429) is not a
/// refusal: the same cookie is checked again and can still sign in.
#[test]
fn a_check_that_fails_in_passing_is_retried() {
let profile = temp_dir("transient");
let browser = FakeBrowser::start(Fake::default());
write_endpoint(&profile, browser.port);
let mut devtools = connect(&profile, &|| true).unwrap();
let checks = std::cell::Cell::new(0);
wait_for_session(
&mut devtools,
Duration::from_secs(5),
&|_| {
checks.set(checks.get() + 1);
match checks.get() {
1 => Err(IdxError::RateLimited),
_ => Ok(()),
}
},
&|| true,
)
.unwrap();
assert_eq!(checks.get(), 2);
browser.finish();
let _ = std::fs::remove_dir_all(&profile);
}
/// Closing the window ends the wait at once, not at the timeout.
#[test]
fn closing_the_window_ends_the_wait() {

View file

@ -72,13 +72,24 @@ impl LoginPage {
///
/// Each connection gets its own thread: browsers open idle preconnects,
/// and answering them one at a time used to hold the page back for the
/// read timeout on every one.
/// read timeout on every one. Pastes are checked one at a time, and none
/// after one is accepted: two at once (a double click, the bookmarklet
/// and the clipboard button) would each store a session, and the later
/// could overwrite the one Stockbit had just rotated.
pub fn wait<F>(&self, timeout: Duration, verify: F) -> Result<(), IdxError>
where
F: Fn(&str) -> Result<(), IdxError> + Send + Sync + 'static,
{
let (done_tx, done_rx) = std::sync::mpsc::channel::<()>();
let verify = std::sync::Arc::new(verify);
let accepted = std::sync::Mutex::new(false);
let verify = std::sync::Arc::new(move |paste: &str| {
let mut accepted = accepted.lock().unwrap_or_else(|e| e.into_inner());
if !*accepted {
verify(paste)?;
*accepted = true;
}
Ok(())
});
let page = std::sync::Arc::new(self.page());
let path = std::sync::Arc::new(format!("/login/{}", self.secret));
let busy = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
@ -575,6 +586,37 @@ mod tests {
(status, body)
}
/// Two pastes at once are checked one at a time, and the second is not
/// checked at all once the first is accepted.
#[test]
fn concurrent_pastes_are_checked_once() {
let page = LoginPage::bind(0).unwrap();
let url = page.url();
let checks = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
let counted = checks.clone();
let handle = std::thread::spawn(move || {
page.wait(Duration::from_secs(10), move |_| {
counted.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
std::thread::sleep(Duration::from_millis(300));
Ok(())
})
});
std::thread::sleep(Duration::from_millis(100));
let clients: Vec<_> = (0..2)
.map(|_| {
let url = url.clone();
std::thread::spawn(move || request(&url, "POST", "good"))
})
.collect();
for client in clients {
let (status, body) = client.join().unwrap();
assert_eq!(status, 200);
assert!(body.contains("\"ok\":true"), "{body}");
}
handle.join().unwrap().unwrap();
assert_eq!(checks.load(std::sync::atomic::Ordering::SeqCst), 1);
}
#[test]
fn serves_the_page_only_with_the_secret_and_accepts_a_good_paste() {
let page = LoginPage::bind(0).unwrap();

View file

@ -252,6 +252,11 @@ pub(crate) fn fundamentals_unsupported(symbol: &str) -> IdxError {
/// Stores a pasted session (stockbit.com's `credentialStorage` cookie or a
/// refresh token) and checks it with one request, which refreshes it when the
/// access token is missing or stale. A session Stockbit refuses is removed.
///
/// A check that fails another way (network, rate limit) puts the previous
/// session back too, unless Stockbit already rotated the pasted one: the
/// refresh spent the pasted token, so the stored pair is then the only
/// working copy, and a refresh Stockbit answered is proof enough.
pub fn login(pasted: &str) -> Result<auth::SessionStatus, IdxError> {
let session = auth::session_from_paste(pasted).ok_or_else(|| {
IdxError::InvalidInput(
@ -264,17 +269,40 @@ pub fn login(pasted: &str) -> Result<auth::SessionStatus, IdxError> {
// A refused paste must not lose a session that already worked.
let previous = std::fs::read(store.path()).ok();
store.save(&session)?;
let saved = std::fs::read(store.path()).ok();
#[cfg(feature = "testing")]
let check = !crate::api::mock::enabled();
#[cfg(not(feature = "testing"))]
let check = true;
if check && let Err(err) = Exodus::new(store.clone()).verify() {
let rotated = std::fs::read(store.path()).ok() != saved;
if !keep_after_failed_check(&err, rotated) {
restore(&store, previous);
return Err(err);
}
crate::runtime::warn(format!(
"Stockbit accepted the session and rotated it, but the check after that failed \
({err}); the new session is stored"
));
}
Ok(store.status(chrono::Utc::now().timestamp()))
}
/// Whether `login` keeps the stored session after its check failed with
/// `err`: only when the check rotated it and Stockbit did not then refuse it.
fn keep_after_failed_check(err: &IdxError, rotated: bool) -> bool {
rotated && !matches!(err, IdxError::AuthError(_))
}
/// Whether a session check failed for a reason that says nothing about the
/// session itself, so the same session may be checked again later.
pub(crate) fn check_is_transient(err: &IdxError) -> bool {
matches!(
err,
IdxError::Http(_) | IdxError::RateLimited | IdxError::UpstreamBlocked(_)
)
}
fn restore(store: &SessionStore, previous: Option<Vec<u8>>) {
match previous {
Some(bytes) => {
@ -301,7 +329,24 @@ pub fn logout() -> Result<(bool, auth::SessionStatus), IdxError> {
#[cfg(test)]
mod tests {
use super::USER_AGENT;
use super::{USER_AGENT, check_is_transient, keep_after_failed_check};
use crate::error::IdxError;
#[test]
fn login_keeps_a_rotated_session_only_when_stockbit_did_not_refuse_it() {
let network = IdxError::Http("timeout".to_string());
let refused = IdxError::AuthError("expired".to_string());
assert!(keep_after_failed_check(&network, true));
assert!(!keep_after_failed_check(&network, false));
assert!(!keep_after_failed_check(&refused, true));
assert!(!keep_after_failed_check(&refused, false));
assert!(check_is_transient(&IdxError::RateLimited));
assert!(check_is_transient(&network));
assert!(!check_is_transient(&refused));
assert!(!check_is_transient(&IdxError::InvalidInput(
"x".to_string()
)));
}
/// CloudFront answers HTTP 403 to agents containing these (probed
/// 2026-10-22); the honest agent must never match one.