session.rs

//! セッション。クッキーには id だけを入れ、中身は置き場に保存する。
//!
//! 置き場は SESSION_STORE で選ぶ（none / memory / file / sqlite）。
//! 表の名前も置き場のディレクトリもクッキー名も、決め打ちにせず使う側が選べる。
//! id は 256 bit の乱数で、判定は置き場にあるかどうかなので署名はしない。
//!
//! 発行は遅らせる。読むだけなら何も作らず、Set-Cookie も出さない。
//! 書いたときに初めて id を作り、保存し、クッキーを出す。
//! こうしないと、初めて来た相手 1 人につき 1 件、クローラの分まで溜まっていく。

use anyhow::Context;
use axum::Router;
use axum::extract::{FromRequestParts, Request, State};
use axum::http::request::Parts;
use axum::http::{HeaderValue, StatusCode, header};
use axum::middleware::Next;
use axum::response::Response;
use axum::routing::get;
use chrono::{DateTime, Utc};
use crate::conf::Conf;
use crate::error::AppError;
use crate::extract::Json;
use crate::state::AppState;
use serde::{Deserialize, Serialize};
use sqlx::SqlitePool;
use std::collections::HashMap;
use std::ops::{Deref, DerefMut};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};


/// 期限切れを掃除する間隔
const SWEEP_INTERVAL: Duration = Duration::from_secs(60 * 60);

/// id の長さ（16 進で書いた 32 バイト）
const ID_LENGTH: usize = 64;


/// セッションに入れておくもの。項目はここに足す。
/// JSON で保存するので、足しても置き場の側は変わらない。
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct SessionData {
    /// 利用者が選んだ言語
    pub language: Option<String>,
}

/// 置き場に入っている 1 件
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Record {
    /// 入れておいたもの
    pub data: SessionData,
    /// 最後に保存した時刻。期限はここから測る
    pub updated_at: DateTime<Utc>,
}

/// 置き場。増やすときはここに 1 つ足すと、扱っていない場所をコンパイラが指す
#[derive(Clone)]
pub enum Store {
    /// セッションを使わない
    None,
    /// プロセスの中。再起動で消える
    Memory(Arc<RwLock<HashMap<String, Record>>>),
    /// ディレクトリの中に `<id>.json` として置く
    File(PathBuf),
    /// SQLite の表に置く。表の名前も使う側が決める
    Sqlite(SqlitePool, String),
}

/// ハンドラが受け取る取っ手。中身は共有なので Clone は安価である
#[derive(Clone)]
pub struct Session {
    inner: Arc<RwLock<Inner>>,
}

struct Inner {
    data: SessionData,
    changed: bool,
    destroyed: bool,
    renewed: bool,
}

/// POST /api/session の本文
#[derive(Debug, Deserialize)]
pub struct WriteSession {
    /// 入れ直す言語
    pub language: Option<String>,
}


/// セッションの表を用意する。sqlite を選んだときだけ効く。
/// 使う側のマイグレーションに混ざらないよう、この表はライブラリが自分で作る。
pub async fn prepare(store: &Store) -> anyhow::Result<()> {
    let Store::Sqlite(database, table) = store else {
        return Ok(());
    };

    sqlx::query(&format!(
        "create table if not exists {table} (\
         id text primary key not null, data text not null, updated_at text not null)"
    ))
    .execute(database)
    .await
    .with_context(|| format!("{table} の表を作れません"))?;

    sqlx::query(&format!(
        "create index if not exists {table}_updated_at on {table} (updated_at)"
    ))
    .execute(database)
    .await
    .with_context(|| format!("{table} の索引を作れません"))?;

    Ok(())
}


/// 設定から置き場を組み立てる
pub fn store(conf: &Conf, database: &SqlitePool) -> Store {
    match conf.session_store.as_str() {
        "memory" => Store::Memory(Arc::new(RwLock::new(HashMap::new()))),
        "file" => Store::File(PathBuf::from(&conf.session_directory)),
        "sqlite" => Store::Sqlite(database.clone(), conf.session_table.clone()),
        _ => Store::None,
    }
}


/// 1 件を読む。期限が切れていたら消して None を返す
pub async fn load(store: &Store, id: &str, max_age: Duration) -> Option<Record> {
    let record = match store {
        Store::None => None,
        Store::Memory(map) => map.read().await.get(id).cloned(),
        Store::File(directory) => read_file(directory, id).await,
        Store::Sqlite(database, table) => read_row(database, table, id).await,
    }?;

    if expired(&record, max_age) {
        forget(store, id).await;
        return None;
    }

    Some(record)
}


/// 1 件を保存する。保存した時刻から期限を測り直す
pub async fn save(store: &Store, id: &str, data: &SessionData) -> anyhow::Result<()> {
    let record = Record {data: data.clone(), updated_at: Utc::now()};

    match store {
        Store::None => {}

        Store::Memory(map) => {
            map.write().await.insert(id.to_owned(), record);
        }

        Store::File(directory) => {
            tokio::fs::create_dir_all(directory)
                .await
                .with_context(|| format!("{} を作れません", directory.display()))?;
            let text = serde_json::to_string(&record)?;
            tokio::fs::write(file_path(directory, id), text)
                .await
                .with_context(|| format!("{id}.json を書けません"))?;
        }

        Store::Sqlite(database, table) => {
            sqlx::query(&format!(
                "insert into {table} (id, data, updated_at) values (?, ?, ?) \
                 on conflict(id) do update set \
                 data = excluded.data, updated_at = excluded.updated_at"
            ))
            .bind(id)
            .bind(serde_json::to_string(&record.data)?)
            .bind(record.updated_at)
            .execute(database)
            .await?;
        }
    }

    Ok(())
}


/// 1 件を消す。消せなくても続ける
pub async fn forget(store: &Store, id: &str) {
    let result: anyhow::Result<()> = async {
        match store {
            Store::None => {}

            Store::Memory(map) => {
                map.write().await.remove(id);
            }

            Store::File(directory) => {
                let path = file_path(directory, id);
                if tokio::fs::try_exists(&path).await.unwrap_or(false) {
                    tokio::fs::remove_file(&path).await?;
                }
            }

            Store::Sqlite(database, table) => {
                sqlx::query(&format!("delete from {table} where id = ?"))
                    .bind(id)
                    .execute(database)
                    .await?;
            }
        }
        Ok(())
    }
    .await;

    if let Err(error) = result {
        tracing::warn!(error = ?error, "セッションを消せません");
    }
}


/// 期限切れを掃除する
pub async fn sweep(store: &Store, max_age: Duration) -> anyhow::Result<()> {
    match store {
        Store::None => {}

        Store::Memory(map) => {
            map.write().await.retain(|_, record| !expired(record, max_age));
        }

        Store::File(directory) => {
            let Ok(mut entries) = tokio::fs::read_dir(directory).await else {
                return Ok(());
            };

            while let Some(entry) = entries.next_entry().await? {
                let Ok(detail) = entry.metadata().await else {
                    continue;
                };
                let stale = detail
                    .modified()
                    .ok()
                    .and_then(|at| at.elapsed().ok())
                    .is_some_and(|age| age > max_age);

                if stale {
                    tokio::fs::remove_file(entry.path()).await?;
                }
            }
        }

        Store::Sqlite(database, table) => {
            sqlx::query(&format!("delete from {table} where updated_at < ?"))
                .bind(limit_time(max_age))
                .execute(database)
                .await?;
        }
    }

    Ok(())
}


/// 起動してすぐ 1 回、そのあとは 1 時間ごとに掃除する
pub fn start_sweep(store: Store, max_age: Duration) {
    tokio::spawn(async move {
        let mut ticker = tokio::time::interval(SWEEP_INTERVAL);

        loop {
            ticker.tick().await;

            if let Err(error) = sweep(&store, max_age).await {
                tracing::warn!(error = ?error, "セッションを掃除できません");
            }
        }
    });
}


/// クッキーを読み、あれば中身を載せてハンドラへ渡す。
/// 書かれていたら保存し、Set-Cookie を付けて返す。
pub async fn handle(State(state): State<AppState>, mut request: Request, next: Next) -> Response {
    let conf = state.conf.clone();

    // 置き場がないなら、読み込みも発行もしない。取っ手だけ渡してハンドラを通す
    if matches!(state.session, Store::None) {
        let session = Session {
            inner: Arc::new(RwLock::new(Inner {
                data: SessionData::default(),
                changed: false,
                destroyed: false,
                renewed: false,
            })),
        };
        request.extensions_mut().insert(session);
        return next.run(request).await;
    }

    let sent = cookie_value(request.headers(), &conf.session_cookie).filter(|value| is_id(value));

    // 覚えのない id は採らない。相手が決めた id をそのまま使うと、あとから乗っ取られる
    let found = match sent {
        Some(id) => load(&state.session, &id, conf.session_max_age)
            .await
            .map(|record| (id, record)),
        None => None,
    };

    let data = found
        .as_ref()
        .map(|(_, record)| record.data.clone())
        .unwrap_or_default();
    let session = Session {
        inner: Arc::new(RwLock::new(Inner {
            data,
            changed: false,
            destroyed: false,
            renewed: false,
        })),
    };
    request.extensions_mut().insert(session.clone());

    let mut response = next.run(request).await;
    let inner = session.inner.read().await;

    if inner.destroyed {
        if let Some((id, _)) = &found {
            forget(&state.session, id).await;
        }
        put_cookie(&mut response, &conf, "", Duration::ZERO);
        return response;
    }

    // 期限の半分を過ぎていたら、書き換えがなくても保存し直して先へ延ばす
    let halfway = found
        .as_ref()
        .is_some_and(|(_, record)| expired(record, max_age_half(conf.session_max_age)));

    // 保存しないなら、クッキーも出さない。読むだけの相手には何も出さず、
    // 変わっていないセッションにも付けない。毎回付けると共有キャッシュに載らなくなる
    if !inner.changed && !halfway {
        return response;
    }

    // 作り直すなら、古いほうを先に捨てる
    if inner.renewed && let Some((id, _)) = &found {
        forget(&state.session, id).await;
    }

    let id = match &found {
        Some((id, _)) if !inner.renewed => id.clone(),
        _ => new_id(),
    };

    if let Err(error) = save(&state.session, &id, &inner.data).await {
        tracing::error!(error = ?error, "セッションを保存できません");
        return response;
    }

    put_cookie(&mut response, &conf, &id, conf.session_max_age);
    response
}


/// セッションを読む・書く・捨てるルート
pub fn router() -> Router<AppState> {
    Router::new().route("/", get(read).post(write).delete(destroy))
}


/// GET /api/session
async fn read(session: Session) -> Json<SessionData> {
    Json(session.read().await.clone())
}


/// POST /api/session  `{"language": "ja"}`
async fn write(session: Session, Json(body): Json<WriteSession>) -> Json<SessionData> {
    let mut data = session.write().await;
    data.language = body.language;
    Json(data.clone())
}


/// DELETE /api/session
async fn destroy(session: Session) -> StatusCode {
    session.destroy().await;
    StatusCode::NO_CONTENT
}


impl Session {
    /// 中身を見る。見るだけでは保存しない
    pub async fn read(&self) -> SessionRead<'_> {
        SessionRead {inner: self.inner.read().await}
    }

    /// 中身を書き換える。呼んだ時点で保存する印が立つ
    pub async fn write(&self) -> SessionWrite<'_> {
        let mut inner = self.inner.write().await;
        inner.changed = true;
        SessionWrite {inner}
    }

    /// 中身は残したまま、id だけ作り直す。
    ///
    /// ログインのように相手の立場が変わったところで呼ぶ。呼ばないと、
    /// 攻撃者が自分で取った id を相手のブラウザに置き、そのまま認証を通せる。
    /// mon は覚えのない id を採らないので id の捏造は防げるが、
    /// 正規に取った id を置かれる筋は、こちらで id を替えないと塞げない。
    pub async fn renew(&self) {
        let mut inner = self.inner.write().await;
        inner.changed = true;
        inner.renewed = true;
    }

    /// 置き場から消し、クッキーも消す
    pub async fn destroy(&self) {
        let mut inner = self.inner.write().await;
        inner.data = SessionData::default();
        inner.destroyed = true;
    }
}


impl<S: Send + Sync> FromRequestParts<S> for Session {
    type Rejection = AppError;

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        parts
            .extensions
            .get::<Session>()
            .cloned()
            .ok_or_else(|| AppError::Internal(anyhow::anyhow!("セッションの層が掛かっていません")))
    }
}


/// 見るための取っ手。[`SessionData`] としてそのまま扱える
pub struct SessionRead<'a> {
    inner: RwLockReadGuard<'a, Inner>,
}

/// 書き換えるための取っ手。作った時点で保存する印が立つ
pub struct SessionWrite<'a> {
    inner: RwLockWriteGuard<'a, Inner>,
}


impl Deref for SessionRead<'_> {
    type Target = SessionData;

    fn deref(&self) -> &SessionData {
        &self.inner.data
    }
}


impl Deref for SessionWrite<'_> {
    type Target = SessionData;

    fn deref(&self) -> &SessionData {
        &self.inner.data
    }
}


impl DerefMut for SessionWrite<'_> {
    fn deref_mut(&mut self) -> &mut SessionData {
        &mut self.inner.data
    }
}


/// 32 バイトの乱数を 16 進で書いたもの
fn new_id() -> String {
    let mut seed = [0u8; 32];
    getrandom::fill(&mut seed).expect("乱数を取れません");
    seed.iter().map(|byte| format!("{byte:02x}")).collect()
}


/// 自分が出した形の id かどうか。ファイル名に使うので、ここを通らないものは触らない
fn is_id(value: &str) -> bool {
    value.len() == ID_LENGTH && value.bytes().all(|byte| byte.is_ascii_hexdigit())
}


fn file_path(directory: &Path, id: &str) -> PathBuf {
    directory.join(format!("{id}.json"))
}


async fn read_file(directory: &Path, id: &str) -> Option<Record> {
    let text = tokio::fs::read_to_string(file_path(directory, id)).await.ok()?;
    serde_json::from_str(&text).ok()
}


async fn read_row(database: &SqlitePool, table: &str, id: &str) -> Option<Record> {
    let row: Option<(String, DateTime<Utc>)> =
        sqlx::query_as(&format!("select data, updated_at from {table} where id = ?"))
            .bind(id)
            .fetch_optional(database)
            .await
            .ok()?;

    let (data, updated_at) = row?;
    Some(Record {data: serde_json::from_str(&data).ok()?, updated_at})
}


fn expired(record: &Record, max_age: Duration) -> bool {
    (Utc::now() - record.updated_at)
        .to_std()
        .is_ok_and(|age| age > max_age)
}


fn max_age_half(max_age: Duration) -> Duration {
    max_age / 2
}


fn limit_time(max_age: Duration) -> DateTime<Utc> {
    Utc::now() - chrono::Duration::from_std(max_age).unwrap_or(chrono::Duration::zero())
}


fn cookie_value(headers: &axum::http::HeaderMap, name: &str) -> Option<String> {
    let header = headers.get(header::COOKIE)?.to_str().ok()?;
    let prefix = format!("{name}=");

    for part in header.split(';') {
        let trimmed = part.trim();

        if let Some(value) = trimmed.strip_prefix(&prefix) {
            return Some(value.to_owned());
        }
    }

    None
}


fn put_cookie(response: &mut Response, conf: &Conf, id: &str, max_age: Duration) {
    let mut line = format!(
        "{}={}; Path=/; HttpOnly; SameSite={}; Max-Age={}",
        conf.session_cookie,
        id,
        same_site(&conf.session_same_site),
        max_age.as_secs()
    );

    if conf.session_secure {
        line.push_str("; Secure");
    }

    if let Ok(value) = HeaderValue::from_str(&line) {
        response.headers_mut().append(header::SET_COOKIE, value);
    }
}


fn same_site(value: &str) -> &'static str {
    match value {
        "strict" => "Strict",
        "none" => "None",
        _ => "Lax",
    }
}