portal/src/sessions.rs

//! Sessions in NATS KV.
//!
//! A session used to live in a `MemoryStore`, which meant every
//! restart signed everybody out: a deploy logged out every desk on
//! every site, and the only reason that was bearable is that the desks
//! were three people. This keeps them in the JetStream KV that is
//! already under everything else here - no new service, no new
//! dependency, and the same bucket a second instance of a site would
//! read if there ever is one.
//!
//! Expiry is said twice on purpose. The bucket carries a `max_age`, so
//! NATS collects abandoned sessions without anyone sweeping; and a
//! record's own `expiry_date` is checked on the way out, because the
//! bucket's age is a coarse fence and the record's date is the truth.
#![cfg(feature = "ssr")]

use tower_sessions::session::{Id, Record};
use tower_sessions::session_store::{Error, Result};

/// How long an untouched session survives in the bucket. Comfortably
/// longer than any expiry `SessionManagerLayer` will set, so the
/// record's own date is always what decides; this is only here so a
/// session nobody ever comes back for does not sit in NATS forever.
const MAX_AGE: std::time::Duration = std::time::Duration::from_secs(60 * 60 * 24 * 40);

pub const BUCKET: &str = "portal_sessions";

/// A new id is drawn this many times before a collision is reported as
/// a backend error. Two random `i128`s colliding is not a thing that
/// happens; this exists so a bucket misbehaving cannot spin forever.
const ID_ATTEMPTS: usize = 8;

#[derive(Clone)]
pub struct NatsStore {
    store: async_nats::jetstream::kv::Store,
}

impl std::fmt::Debug for NatsStore {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("NatsStore")
            .field("bucket", &BUCKET)
            .finish()
    }
}

impl NatsStore {
    /// Opens the bucket, creating it on first use the way every other
    /// bucket here is created.
    pub async fn open(js: &async_nats::jetstream::Context) -> anyhow::Result<Self> {
        let store = match js.get_key_value(BUCKET).await {
            Ok(store) => store,
            Err(_) => {
                js.create_key_value(async_nats::jetstream::kv::Config {
                    bucket: BUCKET.to_string(),
                    max_age: MAX_AGE,
                    history: 1,
                    ..Default::default()
                })
                .await?
            }
        };
        Ok(Self { store })
    }

    fn encode(record: &Record) -> Result<Vec<u8>> {
        serde_json::to_vec(record).map_err(|e| Error::Encode(e.to_string()))
    }
}

#[async_trait::async_trait]
impl tower_sessions::SessionStore for NatsStore {
    /// `create` draws a fresh id until the bucket accepts one as new.
    /// `put` would overwrite a stranger's session on a collision;
    /// `create` refuses, which is the whole difference.
    async fn create(&self, record: &mut Record) -> Result<()> {
        for _ in 0..ID_ATTEMPTS {
            let bytes = Self::encode(record)?;
            match self
                .store
                .create(&record.id.to_string(), bytes.into())
                .await
            {
                Ok(_) => return Ok(()),
                Err(_) => record.id = Id::default(),
            }
        }
        Err(Error::Backend(format!(
            "could not find an unused session id in {ID_ATTEMPTS} attempts"
        )))
    }

    async fn save(&self, record: &Record) -> Result<()> {
        let bytes = Self::encode(record)?;
        self.store
            .put(&record.id.to_string(), bytes.into())
            .await
            .map(|_| ())
            .map_err(|e| Error::Backend(e.to_string()))
    }

    async fn load(&self, id: &Id) -> Result<Option<Record>> {
        let Some(bytes) = self
            .store
            .get(&id.to_string())
            .await
            .map_err(|e| Error::Backend(e.to_string()))?
        else {
            return Ok(None);
        };
        let record: Record =
            serde_json::from_slice(&bytes).map_err(|e| Error::Decode(e.to_string()))?;
        // An expired record is gone as far as anyone asking is
        // concerned, whether or not the bucket has collected it yet.
        if record.expiry_date <= tower_sessions::cookie::time::OffsetDateTime::now_utc() {
            return Ok(None);
        }
        Ok(Some(record))
    }

    async fn delete(&self, id: &Id) -> Result<()> {
        // Purge, not delete: a signed-out session should leave no
        // history in the bucket to read back.
        self.store
            .purge(&id.to_string())
            .await
            .map_err(|e| Error::Backend(e.to_string()))
    }
}