gyazo-mcp-server 0.6.3

use std::{collections::HashMap, fs, path::Path};

use anyhow::{Context, Result};
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use hmac::{Hmac, Mac};
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use uuid::Uuid;

use crate::{
    app_state::{AccessTokenRecord, RegisteredClient},
    gyazo_api::GyazoUserProfile,
};

type HmacSha256 = Hmac<Sha256>;

const TOKEN_PREFIX: &str = "gmcp1";

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) struct StoredMcpSessionState {
    pub(crate) signing_key: String,
    #[serde(default)]
    pub(crate) sessions: Vec<StoredMcpSession>,
    /// 既存ファイルとの後方互換のため、`#[serde(default)]` で空ベクタを許容する。
    #[serde(default)]
    pub(crate) registered_clients: Vec<StoredRegisteredClient>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) struct StoredMcpSession {
    pub(crate) session_id: String,
    pub(crate) backend_access_token: String,
    pub(crate) gyazo_user: GyazoUserProfile,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) struct StoredRegisteredClient {
    pub(crate) client_id: String,
    pub(crate) redirect_uris: Vec<String>,
}

pub(crate) fn load_mcp_session_state(path: &Path) -> Result<Option<StoredMcpSessionState>> {
    if !path.exists() {
        return Ok(None);
    }

    let raw = fs::read_to_string(path).with_context(|| {
        format!(
            "MCP session file を読み取れませんでした: {}",
            path.display()
        )
    })?;
    let state = toml::from_str(&raw).with_context(|| {
        format!(
            "MCP session file を解析できませんでした: {}",
            path.display()
        )
    })?;

    Ok(Some(state))
}

pub(crate) fn save_mcp_session_state(path: &Path, state: &StoredMcpSessionState) -> Result<()> {
    let raw = toml::to_string(state).context("MCP session file をシリアライズできませんでした")?;

    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent).with_context(|| {
            format!(
                "MCP session directory を作成できませんでした: {}",
                parent.display()
            )
        })?;
    }

    fs::write(path, raw).with_context(|| {
        format!(
            "MCP session file に書き込めませんでした: {}",
            path.display()
        )
    })
}

pub(crate) fn generate_signing_key() -> Vec<u8> {
    let mut signing_key = Vec::with_capacity(32);
    signing_key.extend_from_slice(Uuid::new_v4().as_bytes());
    signing_key.extend_from_slice(Uuid::new_v4().as_bytes());
    signing_key
}

pub(crate) fn encode_signing_key(signing_key: &[u8]) -> String {
    URL_SAFE_NO_PAD.encode(signing_key)
}

pub(crate) fn decode_signing_key(encoded: &str) -> Result<Vec<u8>> {
    URL_SAFE_NO_PAD
        .decode(encoded)
        .context("MCP signing key をデコードできませんでした")
}

pub(crate) fn sign_access_token(signing_key: &[u8], session_id: &str) -> Result<String> {
    let signature = token_signature(signing_key, session_id)?;
    Ok(format!("{TOKEN_PREFIX}.{session_id}.{signature}"))
}

pub(crate) fn verify_access_token(signing_key: &[u8], token: &str) -> Result<Option<String>> {
    let mut parts = token.split('.');
    let prefix = parts.next();
    let session_id = parts.next();
    let signature = parts.next();

    if prefix != Some(TOKEN_PREFIX) || parts.next().is_some() {
        return Ok(None);
    }

    let Some(session_id) = session_id else {
        return Ok(None);
    };
    let Some(signature) = signature else {
        return Ok(None);
    };

    let expected_signature = token_signature(signing_key, session_id)?;
    if expected_signature != signature {
        return Ok(None);
    }

    Ok(Some(session_id.to_string()))
}

pub(crate) fn sessions_to_records(
    sessions: Vec<StoredMcpSession>,
) -> HashMap<String, AccessTokenRecord> {
    sessions
        .into_iter()
        .map(|session| {
            (
                session.session_id,
                AccessTokenRecord {
                    backend_access_token: session.backend_access_token,
                    gyazo_user: session.gyazo_user,
                },
            )
        })
        .collect()
}

pub(crate) fn records_to_sessions(
    records: &HashMap<String, AccessTokenRecord>,
) -> Vec<StoredMcpSession> {
    records
        .iter()
        .map(|(session_id, record)| StoredMcpSession {
            session_id: session_id.clone(),
            backend_access_token: record.backend_access_token.clone(),
            gyazo_user: record.gyazo_user.clone(),
        })
        .collect()
}

pub(crate) fn stored_clients_to_map(
    clients: Vec<StoredRegisteredClient>,
) -> HashMap<String, RegisteredClient> {
    clients
        .into_iter()
        .map(|client| {
            (
                client.client_id,
                RegisteredClient {
                    redirect_uris: client.redirect_uris,
                },
            )
        })
        .collect()
}

pub(crate) fn map_to_stored_clients(
    clients: &HashMap<String, RegisteredClient>,
) -> Vec<StoredRegisteredClient> {
    clients
        .iter()
        .map(|(client_id, client)| StoredRegisteredClient {
            client_id: client_id.clone(),
            redirect_uris: client.redirect_uris.clone(),
        })
        .collect()
}

fn token_signature(signing_key: &[u8], session_id: &str) -> Result<String> {
    let mut mac =
        HmacSha256::new_from_slice(signing_key).context("HMAC signer を初期化できませんでした")?;
    mac.update(TOKEN_PREFIX.as_bytes());
    mac.update(b":");
    mac.update(session_id.as_bytes());
    let signature = mac.finalize().into_bytes();

    Ok(URL_SAFE_NO_PAD.encode(signature))
}

#[cfg(test)]
mod tests {
    use std::{
        collections::HashMap,
        fs,
        time::{SystemTime, UNIX_EPOCH},
    };

    use super::{
        StoredMcpSessionState, decode_signing_key, encode_signing_key, generate_signing_key,
        load_mcp_session_state, map_to_stored_clients, records_to_sessions, save_mcp_session_state,
        sessions_to_records, sign_access_token, stored_clients_to_map, verify_access_token,
    };
    use crate::{
        app_state::{AccessTokenRecord, RegisteredClient},
        gyazo_api::GyazoUserProfile,
    };

    #[test]
    fn signs_and_verifies_access_token() {
        let signing_key = generate_signing_key();
        let token = sign_access_token(&signing_key, "session-123").unwrap();
        let verified = verify_access_token(&signing_key, &token).unwrap();

        assert_eq!(verified.as_deref(), Some("session-123"));
        assert!(
            verify_access_token(&signing_key, "gmcp1.session-123.invalid")
                .unwrap()
                .is_none()
        );
    }

    #[test]
    fn saves_and_loads_mcp_session_state() {
        let unique = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let dir = std::env::temp_dir().join(format!("gyazo-mcp-sessions-test-{unique}"));
        let path = dir.join("mcp_sessions.toml");
        let signing_key = generate_signing_key();
        let mut records = HashMap::new();
        records.insert(
            "session-1".to_string(),
            AccessTokenRecord {
                backend_access_token: "backend-token".to_string(),
                gyazo_user: GyazoUserProfile {
                    email: "test@example.com".to_string(),
                    name: "tester".to_string(),
                    profile_image: "https://example.com/avatar.png".to_string(),
                    uid: "user-1".to_string(),
                },
            },
        );
        let state = StoredMcpSessionState {
            signing_key: encode_signing_key(&signing_key),
            sessions: records_to_sessions(&records),
            registered_clients: Vec::new(),
        };

        fs::create_dir_all(&dir).unwrap();
        save_mcp_session_state(&path, &state).unwrap();
        let loaded = load_mcp_session_state(&path).unwrap().unwrap();

        assert_eq!(
            decode_signing_key(&loaded.signing_key).unwrap(),
            signing_key
        );
        assert_eq!(sessions_to_records(loaded.sessions), records);

        let _ = fs::remove_file(&path);
        let _ = fs::remove_dir(&dir);
    }

    #[test]
    fn saves_and_loads_registered_clients() {
        let unique = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let dir = std::env::temp_dir().join(format!("gyazo-mcp-clients-test-{unique}"));
        let path = dir.join("mcp_sessions.toml");
        let signing_key = generate_signing_key();

        let mut clients = HashMap::new();
        clients.insert(
            "client-abc".to_string(),
            RegisteredClient {
                redirect_uris: vec![
                    "http://127.0.0.1:18449/oauth/callback".to_string(),
                    "http://localhost:18449/oauth/callback".to_string(),
                ],
            },
        );
        clients.insert(
            "client-xyz".to_string(),
            RegisteredClient {
                redirect_uris: vec!["http://127.0.0.1:18449/cb".to_string()],
            },
        );

        let state = StoredMcpSessionState {
            signing_key: encode_signing_key(&signing_key),
            sessions: Vec::new(),
            registered_clients: map_to_stored_clients(&clients),
        };

        fs::create_dir_all(&dir).unwrap();
        save_mcp_session_state(&path, &state).unwrap();
        let loaded = load_mcp_session_state(&path).unwrap().unwrap();

        let restored = stored_clients_to_map(loaded.registered_clients);
        assert_eq!(restored.len(), 2);
        assert_eq!(
            restored.get("client-abc").unwrap().redirect_uris,
            vec![
                "http://127.0.0.1:18449/oauth/callback".to_string(),
                "http://localhost:18449/oauth/callback".to_string(),
            ]
        );
        assert_eq!(
            restored.get("client-xyz").unwrap().redirect_uris,
            vec!["http://127.0.0.1:18449/cb".to_string()]
        );

        let _ = fs::remove_file(&path);
        let _ = fs::remove_dir(&dir);
    }

    #[test]
    fn loads_legacy_session_file_without_registered_clients_field() {
        // 既存の mcp_sessions.toml は registered_clients フィールドを持たないため、
        // 後方互換のために `#[serde(default)]` で空ベクタとして読めることを保証する。
        let unique = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let dir = std::env::temp_dir().join(format!("gyazo-mcp-legacy-test-{unique}"));
        let path = dir.join("mcp_sessions.toml");
        fs::create_dir_all(&dir).unwrap();

        // フィールドを意図的に省略した legacy 形式の TOML
        fs::write(
            &path,
            "signing_key = \"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA\"\n",
        )
        .unwrap();

        let loaded = load_mcp_session_state(&path).unwrap().unwrap();
        assert!(loaded.sessions.is_empty());
        assert!(loaded.registered_clients.is_empty());

        let _ = fs::remove_file(&path);
        let _ = fs::remove_dir(&dir);
    }
}