nkode-protocol/src/server.rs

use opaque_ke::{
    rand::rngs::OsRng, CredentialFinalization, CredentialRequest,
    RegistrationRequest, ServerLogin, ServerLoginParameters,
    ServerRegistration,
};
use uuid::Uuid;
use crate::models::{RegisterSession, LoginSession, NKodeCipherSuite, NKodeServerSetup, PasswordFile};

#[derive(Debug)]
pub enum AuthRepoError {
    UserExists,
    KeyNotRegistered,
    CodeNotRegistered,
}

pub trait AuthRepo {
    fn new_key(&mut self, identifier: &[u8], password_file: PasswordFile) -> Result<(), AuthRepoError>;
    fn new_code(&mut self, identifier: &[u8], password_file: PasswordFile) -> Result<(), AuthRepoError>;

    fn has_code(&self, identifier: &[u8]) -> bool;
    fn has_key(&self, identifier: &[u8]) -> bool;

    fn get_key_passcode_file(&self, identifier: &[u8]) -> Result<PasswordFile, AuthRepoError>;
    fn get_code_passcode_file(&self, identifier: &[u8]) -> Result<PasswordFile, AuthRepoError>;
}

pub struct RegCache {
    pub session_id: Uuid,
    pub identifier: Vec<u8>,
}

pub struct LoginCache {
    pub session_id: Uuid,
    pub identifiers: Vec<u8>,
    pub server_login: ServerLogin<NKodeCipherSuite>,
}

pub trait AuthSession {
    fn new_reg_session(&mut self, identifier: &[u8]) -> Result<RegCache, String>;
    fn get_reg_session(&self, session_id: &Uuid) -> Result<RegCache, String>;
    fn clear_reg_session(&mut self, session_id: &Uuid) -> Result<(), String>;

    fn new_login_session(
        &mut self,
        identifier: &[u8],
        server_login: ServerLogin<NKodeCipherSuite>,
    ) -> Result<LoginCache, String>;
    fn get_login_session(&self, session_id: &Uuid) -> Result<LoginCache, String>;
    fn clear_login_session(&mut self, session_id: &Uuid) -> Result<(), String>;
}


pub trait CredKind {
    fn has<R: AuthRepo>(repo: &R, id: &[u8]) -> bool;
    fn get_pf<R: AuthRepo>(repo: &R, id: &[u8]) -> Result<PasswordFile, AuthRepoError>;
    fn put_pf<R: AuthRepo>(repo: &mut R, id: &[u8], pf: PasswordFile) -> Result<(), AuthRepoError>;
    fn prereq_for_register<R: AuthRepo>(_repo: &R, _id: &[u8]) -> Result<(), AuthRepoError> {
        Ok(())
    }
}

pub struct Key;
pub struct Code;

impl CredKind for Key {
    fn has<R: AuthRepo>(repo: &R, id: &[u8]) -> bool {
        repo.has_key(id)
    }
    fn get_pf<R: AuthRepo>(repo: &R, id: &[u8]) -> Result<PasswordFile, AuthRepoError> {
        repo.get_key_passcode_file(id)
    }
    fn put_pf<R: AuthRepo>(repo: &mut R, id: &[u8], pf: PasswordFile) -> Result<(), AuthRepoError> {
        repo.new_key(id, pf)
    }
}

impl CredKind for Code {
    fn has<R: AuthRepo>(repo: &R, id: &[u8]) -> bool {
        repo.has_code(id)
    }
    fn get_pf<R: AuthRepo>(repo: &R, id: &[u8]) -> Result<PasswordFile, AuthRepoError> {
        repo.get_code_passcode_file(id)
    }
    fn put_pf<R: AuthRepo>(repo: &mut R, id: &[u8], pf: PasswordFile) -> Result<(), AuthRepoError> {
        repo.new_code(id, pf)
    }
    fn prereq_for_register<R: AuthRepo>(repo: &R, id: &[u8]) -> Result<(), AuthRepoError> {
        if repo.has_key(id) {
            Ok(())
        } else {
            Err(AuthRepoError::KeyNotRegistered)
        }
    }
}

pub struct OpaqueAuth<R: AuthRepo, S: AuthSession> {
    server_setup: NKodeServerSetup,
    user_repo: R,
    session: S,
}

impl<R: AuthRepo, S: AuthSession> OpaqueAuth<R, S> {
    pub fn new(server_setup: NKodeServerSetup, user_repo: R, session: S) -> Self {
        Self { server_setup, user_repo, session }
    }

    pub async fn reg_start<K: CredKind>(
        &mut self,
        identifier: &[u8],
        request: RegistrationRequest<NKodeCipherSuite>,
    ) -> Result<RegisterSession, String> {
        K::prereq_for_register(&self.user_repo, identifier)
            .map_err(|e| format!("registration prereq failed: {e:?}"))?;
        let start = ServerRegistration::<NKodeCipherSuite>::start(
            &self.server_setup,
            request,
            identifier,
        ).map_err(|e| format!("opaque reg start: {e:?}"))?;
        let cache = self.session
            .new_reg_session(identifier)
            .map_err(|e| format!("reg cache: {e}"))?;

        Ok(RegisterSession { session_id: cache.session_id, response: start.message })
    }

    pub async fn reg_finish<K: CredKind>(
        &mut self,
        session_id: &Uuid,
        password_file: PasswordFile,
    ) -> Result<(), String> {
        let sess = self.session
            .get_reg_session(session_id)
            .map_err(|e| format!("get reg session: {e}"))?;
        K::prereq_for_register(&self.user_repo, sess.identifier.as_slice())
            .map_err(|e| format!("registration prereq failed: {e:?}"))?;
        K::put_pf(&mut self.user_repo, sess.identifier.as_slice(), password_file)
            .map_err(|e| format!("repo write: {e:?}"))?;
        self.session
            .clear_reg_session(session_id)
            .map_err(|e| format!("clear reg session: {e}"))
    }

    pub async fn login_start<K: CredKind>(
        &mut self,
        identifier: &[u8],
        request: CredentialRequest<NKodeCipherSuite>,
    ) -> Result<LoginSession, String> {
        let password_file = K::get_pf(&self.user_repo, identifier)
            .map_err(|e| format!("repo read: {e:?}"))?;

        let password_file =
            ServerRegistration::<NKodeCipherSuite>::deserialize(password_file.as_slice())
                .map_err(|e| format!("pf deserialize: {e:?}"))?;

        let mut server_rng = OsRng;
        let start = ServerLogin::start(
            &mut server_rng,
            &self.server_setup,
            Some(password_file),
            request,
            identifier,
            ServerLoginParameters::default(),
        ).map_err(|e| format!("opaque login start: {e:?}"))?;

        let cache = self.session
            .new_login_session(identifier, start.state)
            .map_err(|e| format!("login cache: {e}"))?;

        Ok(LoginSession { session_id: cache.session_id, response: start.message })
    }

    pub async fn login_finish<K: CredKind>(
        &mut self,
        session_id: &Uuid,
        finalize: CredentialFinalization<NKodeCipherSuite>,
    ) -> Result<Vec<u8>, String> {
        let cache = self.session
            .get_login_session(session_id)
            .map_err(|e| format!("get login session: {e}"))?;

        let finish = cache.server_login
            .finish(finalize, ServerLoginParameters::default())
            .map_err(|e| format!("opaque login finish: {e:?}"))?;

        self.session
            .clear_login_session(session_id)
            .map_err(|e| format!("clear login session: {e}"))?;

        Ok(finish.session_key.to_vec())
    }
}

// ---------------- Usage notes ----------------
//
// You now have these concrete “types” for your call sites:
//
//   KeyAuthRegistration<R, S>  == OpaqueAuth<Registration<Key>,  R, S>
//   CodeAuthRegistration<R, S> == OpaqueAuth<Registration<Code>, R, S>
//   KeyAuthLogin<R, S>         == OpaqueAuth<Login<Key>,         R, S>
//   CodeAuthLogin<R, S>        == OpaqueAuth<Login<Code>,        R, S>
//
// And you only wrote the reg/login OPAQUE logic once.