khadhroony-bot3/ks-wallet/src/wallet.rs

// file: ks-wallet/src/wallet.rs
// version: 12

//! Local wallet storage and signing primitives.

use solana_signer::Signer; // rust-rules: trait-import
use zeroize::Zeroize; // rust-rules: trait-import

/// Validated non-secret wallet alias used as a local filename stem.
#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub struct WalletAlias(std::string::String);

impl crate::WalletAlias {
    /// Parses and validates a wallet alias.
    pub fn parse(value: impl std::convert::Into<std::string::String>) -> ks_core::Result<Self> {
        let value = value.into();
        if value.is_empty() || value.len() > 64 {
            return std::result::Result::Err(ks_core::Error::new(
                "wallet_alias_length_invalid",
                "wallet alias length must be between 1 and 64 bytes",
            ));
        }
        if !value.bytes().enumerate().all(|(index, byte)| {
            return byte.is_ascii_alphanumeric() || (index > 0 && (byte == b'_' || byte == b'-'));
        }) {
            return std::result::Result::Err(ks_core::Error::new(
                "wallet_alias_invalid",
                "wallet alias must start with an ASCII letter or digit and contain only ASCII letters, digits, '_' or '-'",
            ));
        }
        return std::result::Result::Ok(Self(value));
    }

    /// Returns the validated alias text.
    pub fn as_str(&self) -> &str {
        return self.0.as_str();
    }
}

/// Non-secret wallet policy used by callers before requesting a signature.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct WalletPolicy {
    /// Whether signing is enabled for this wallet.
    pub signing_enabled: bool,
    /// Optional lamport spend limit enforced by the execution layer.
    pub lamport_spend_limit: std::option::Option<u64>,
}

/// Whether a runtime wallet is temporary or persistent.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum WalletPersistence {
    /// Wallet exists only in runtime memory.
    Temporary,
    /// Wallet originates from persistent local storage.
    Persistent,
}

/// Minimal non-secret identity of a wallet.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct WalletIdentity {
    /// Validated logical alias.
    pub alias: crate::WalletAlias,
    /// Base58 Solana public key.
    pub public_key: std::string::String,
    /// Whether the runtime wallet is temporary or persistent.
    pub persistence: crate::WalletPersistence,
}

/// Backward-compatible name for the non-secret wallet identity.
pub type WalletSummary = crate::WalletIdentity;

/// Solana keypair kept private inside the wallet boundary.
pub struct TemporaryWallet {
    alias: crate::WalletAlias,
    keypair: solana_keypair::Keypair,
    storage_path: std::option::Option<std::path::PathBuf>,
}

impl std::fmt::Debug for crate::TemporaryWallet {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        return formatter
            .debug_struct("TemporaryWallet")
            .field("alias", &self.alias)
            .field("public_key", &self.public_key())
            .field("persistence", &self.persistence())
            .finish();
    }
}

impl crate::TemporaryWallet {
    /// Generates a non-persistent wallet in memory.
    pub fn generate(alias: crate::WalletAlias) -> Self {
        return Self {
            alias,
            keypair: solana_keypair::Keypair::new(),
            storage_path: std::option::Option::None,
        };
    }

    /// Returns the minimal non-secret wallet identity.
    pub fn identity(&self) -> crate::WalletIdentity {
        return crate::WalletIdentity {
            alias: self.alias.clone(),
            public_key: self.public_key(),
            persistence: self.persistence(),
        };
    }

    /// Returns the non-secret wallet identity using the historical API name.
    pub fn summary(&self) -> crate::WalletSummary {
        return self.identity();
    }

    /// Returns whether this runtime wallet is temporary or persistent.
    pub fn persistence(&self) -> crate::WalletPersistence {
        if self.storage_path.is_some() {
            return crate::WalletPersistence::Persistent;
        }
        return crate::WalletPersistence::Temporary;
    }

    /// Returns the wallet public key in base58 form.
    pub fn public_key(&self) -> std::string::String {
        return self.keypair.pubkey().to_string();
    }

    /// Returns the signer interface without exposing keypair bytes.
    pub fn as_signer(&self) -> &dyn solana_signer::Signer {
        return &self.keypair;
    }

    /// Returns the signer interface with the `Sync` auto-trait preserved for `Send` async futures.
    pub fn as_sync_signer(&self) -> &(dyn solana_signer::Signer + std::marker::Sync) {
        return &self.keypair;
    }

    /// Signs arbitrary message bytes and returns the base58 signature.
    pub fn sign_message(&self, message: &[u8]) -> ks_core::Result<std::string::String> {
        let signature = match self.keypair.try_sign_message(message) {
            std::result::Result::Ok(signature) => signature,
            std::result::Result::Err(error) => {
                return std::result::Result::Err(ks_core::Error::new(
                    "wallet_message_sign_failed",
                    error.to_string(),
                ));
            },
        };
        tracing::debug!(
            target: crate::TRACING_TARGET,
            action = "sign_message",
            wallet_alias = self.alias.as_str(),
            public_key = %self.keypair.pubkey(),
            message_length = message.len(),
            "signed message with local wallet"
        );
        return std::result::Result::Ok(signature.to_string());
    }
}

/// Filesystem-backed store for development and integration-test wallets.
#[derive(Clone, Eq, PartialEq)]
pub struct TemporaryWalletStore {
    directory: std::path::PathBuf,
}

impl std::fmt::Debug for crate::TemporaryWalletStore {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        return formatter.debug_struct("TemporaryWalletStore").finish();
    }
}

impl crate::TemporaryWalletStore {
    /// Creates a wallet store rooted at the supplied directory.
    pub fn new(directory: impl std::convert::Into<std::path::PathBuf>) -> ks_core::Result<Self> {
        let directory = directory.into();
        if directory.as_os_str().is_empty() {
            return std::result::Result::Err(ks_core::Error::new(
                "wallet_directory_empty",
                "wallet directory must not be empty",
            ));
        }
        return std::result::Result::Ok(Self { directory });
    }

    /// Returns the configured wallet directory.
    pub fn directory(&self) -> &std::path::Path {
        return self.directory.as_path();
    }

    /// Returns the deterministic JSON path for an alias.
    pub fn wallet_path(&self, alias: &crate::WalletAlias) -> std::path::PathBuf {
        return self.directory.join(format!("{}.json", alias.as_str()));
    }

    /// Returns whether a persistent wallet exists for an alias.
    pub async fn exists(&self, alias: &crate::WalletAlias) -> ks_core::Result<bool> {
        let path = self.wallet_path(alias);
        return match tokio::fs::try_exists(&path).await {
            std::result::Result::Ok(exists) => std::result::Result::Ok(exists),
            std::result::Result::Err(error) => std::result::Result::Err(ks_core::Error::new(
                "wallet_file_exists_check_failed",
                error.to_string(),
            )),
        };
    }

    /// Creates and persists a new wallet without overwriting an existing file.
    pub async fn create(
        &self,
        alias: crate::WalletAlias,
    ) -> ks_core::Result<crate::TemporaryWallet> {
        match prepare_wallet_directory(&self.directory).await {
            std::result::Result::Ok(()) => {},
            std::result::Result::Err(error) => return std::result::Result::Err(error),
        }
        let path = self.wallet_path(&alias);
        let keypair = solana_keypair::Keypair::new();
        match write_new_keypair(path.clone(), &keypair).await {
            std::result::Result::Ok(()) => {},
            std::result::Result::Err(error) => return std::result::Result::Err(error),
        }
        tracing::info!(
            target: crate::TRACING_TARGET,
            action = "create_temporary_wallet",
            wallet_alias = alias.as_str(),
            public_key = %keypair.pubkey(),
            "created persistent temporary wallet"
        );
        return std::result::Result::Ok(crate::TemporaryWallet {
            alias,
            keypair,
            storage_path: std::option::Option::Some(path),
        });
    }

    /// Loads and validates an existing persistent wallet.
    pub async fn load(&self, alias: crate::WalletAlias) -> ks_core::Result<crate::TemporaryWallet> {
        let path = self.wallet_path(&alias);
        let keypair = match read_keypair(&path).await {
            std::result::Result::Ok(keypair) => keypair,
            std::result::Result::Err(error) => return std::result::Result::Err(error),
        };
        tracing::debug!(
            target: crate::TRACING_TARGET,
            action = "load_temporary_wallet",
            wallet_alias = alias.as_str(),
            public_key = %keypair.pubkey(),
            "loaded persistent temporary wallet"
        );
        return std::result::Result::Ok(crate::TemporaryWallet {
            alias,
            keypair,
            storage_path: std::option::Option::Some(path),
        });
    }

    /// Loads an existing wallet or creates it when absent.
    pub async fn load_or_create(
        &self,
        alias: crate::WalletAlias,
    ) -> ks_core::Result<crate::TemporaryWallet> {
        let exists = match self.exists(&alias).await {
            std::result::Result::Ok(exists) => exists,
            std::result::Result::Err(error) => return std::result::Result::Err(error),
        };
        if exists {
            return self.load(alias).await;
        }
        let create_result = self.create(alias.clone()).await;
        return match create_result {
            std::result::Result::Ok(wallet) => std::result::Result::Ok(wallet),
            std::result::Result::Err(error) if error.code() == "wallet_file_already_exists" => {
                return self.load(alias).await;
            },
            std::result::Result::Err(error) => std::result::Result::Err(error),
        };
    }
}

async fn prepare_wallet_directory(directory: &std::path::Path) -> ks_core::Result<()> {
    match tokio::fs::create_dir_all(directory).await {
        std::result::Result::Ok(()) => {},
        std::result::Result::Err(error) => {
            return std::result::Result::Err(ks_core::Error::new(
                "wallet_directory_create_failed",
                error.to_string(),
            ));
        },
    }
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt; // rust-rules: trait-import
        let permissions = std::fs::Permissions::from_mode(0o700);
        match tokio::fs::set_permissions(directory, permissions).await {
            std::result::Result::Ok(()) => {},
            std::result::Result::Err(error) => {
                return std::result::Result::Err(ks_core::Error::new(
                    "wallet_directory_permissions_failed",
                    error.to_string(),
                ));
            },
        }
    }
    return std::result::Result::Ok(());
}

async fn write_new_keypair(
    path: std::path::PathBuf,
    keypair: &solana_keypair::Keypair,
) -> ks_core::Result<()> {
    let mut keypair_bytes = keypair.to_bytes();
    let mut encoded = match serde_json::to_vec(keypair_bytes.as_slice()) {
        std::result::Result::Ok(encoded) => encoded,
        std::result::Result::Err(error) => {
            keypair_bytes.zeroize();
            return std::result::Result::Err(ks_core::Error::new(
                "wallet_keypair_serialize_failed",
                error.to_string(),
            ));
        },
    };
    keypair_bytes.zeroize();
    let task_path = path.clone();
    let task_result = tokio::task::spawn_blocking(move || {
        use std::io::Write; // rust-rules: trait-import
        let mut options = std::fs::OpenOptions::new();
        options.write(true).create_new(true);
        #[cfg(unix)]
        {
            use std::os::unix::fs::OpenOptionsExt; // rust-rules: trait-import
            options.mode(0o600);
        }
        let mut file = match options.open(&task_path) {
            std::result::Result::Ok(file) => file,
            std::result::Result::Err(error) => {
                encoded.zeroize();
                let code = if error.kind() == std::io::ErrorKind::AlreadyExists {
                    "wallet_file_already_exists"
                } else {
                    "wallet_file_create_failed"
                };
                return std::result::Result::Err(ks_core::Error::new(code, error.to_string()));
            },
        };
        if let std::result::Result::Err(error) = file.write_all(encoded.as_slice()) {
            encoded.zeroize();
            let _ = std::fs::remove_file(&task_path);
            return std::result::Result::Err(ks_core::Error::new(
                "wallet_file_write_failed",
                error.to_string(),
            ));
        }
        encoded.zeroize();
        if let std::result::Result::Err(error) = file.sync_all() {
            let _ = std::fs::remove_file(&task_path);
            return std::result::Result::Err(ks_core::Error::new(
                "wallet_file_sync_failed",
                error.to_string(),
            ));
        }
        return std::result::Result::Ok(());
    })
    .await;
    return match task_result {
        std::result::Result::Ok(result) => result,
        std::result::Result::Err(error) => std::result::Result::Err(ks_core::Error::new(
            "wallet_file_task_failed",
            error.to_string(),
        )),
    };
}

async fn read_keypair(path: &std::path::Path) -> ks_core::Result<solana_keypair::Keypair> {
    match validate_wallet_file_metadata(path).await {
        std::result::Result::Ok(()) => {},
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    }
    let mut encoded = match tokio::fs::read(path).await {
        std::result::Result::Ok(encoded) => encoded,
        std::result::Result::Err(error) => {
            return std::result::Result::Err(ks_core::Error::new(
                "wallet_file_read_failed",
                error.to_string(),
            ));
        },
    };
    let parsed_result = serde_json::from_slice::<std::vec::Vec<u8>>(encoded.as_slice());
    encoded.zeroize();
    let mut keypair_bytes = match parsed_result {
        std::result::Result::Ok(keypair_bytes) => keypair_bytes,
        std::result::Result::Err(error) => {
            return std::result::Result::Err(ks_core::Error::new(
                "wallet_keypair_json_invalid",
                error.to_string(),
            ));
        },
    };
    if keypair_bytes.len() != crate::SOLANA_KEYPAIR_LENGTH {
        let length = keypair_bytes.len();
        keypair_bytes.zeroize();
        return std::result::Result::Err(ks_core::Error::new(
            "wallet_keypair_length_invalid",
            format!(
                "wallet keypair contains {length} bytes instead of {}",
                crate::SOLANA_KEYPAIR_LENGTH
            ),
        ));
    }
    let keypair_result = solana_keypair::Keypair::try_from(keypair_bytes.as_slice());
    keypair_bytes.zeroize();
    return match keypair_result {
        std::result::Result::Ok(keypair) => std::result::Result::Ok(keypair),
        std::result::Result::Err(error) => std::result::Result::Err(ks_core::Error::new(
            "wallet_keypair_invalid",
            error.to_string(),
        )),
    };
}

async fn validate_wallet_file_metadata(path: &std::path::Path) -> ks_core::Result<()> {
    let metadata = match tokio::fs::symlink_metadata(path).await {
        std::result::Result::Ok(metadata) => metadata,
        std::result::Result::Err(error) => {
            return std::result::Result::Err(ks_core::Error::new(
                "wallet_file_metadata_failed",
                error.to_string(),
            ));
        },
    };
    if metadata.file_type().is_symlink() || !metadata.is_file() {
        return std::result::Result::Err(ks_core::Error::new(
            "wallet_file_type_invalid",
            "wallet file must be a regular file and not a symlink",
        ));
    }
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt; // rust-rules: trait-import
        let mode = metadata.permissions().mode() & 0o777;
        if mode & 0o077 != 0 {
            return std::result::Result::Err(ks_core::Error::new(
                "wallet_file_permissions_too_open",
                format!("wallet file has mode {mode:o}; expected no group or other access"),
            ));
        }
    }
    return std::result::Result::Ok(());
}

#[cfg(test)]
mod tests {
    #[test]
    fn wallet_alias_rejects_path_traversal_and_invalid_characters() {
        for value in ["", ".hidden", "../wallet", "wallet/name", "wallet name"] {
            assert!(crate::WalletAlias::parse(value).is_err());
        }
        assert!(crate::WalletAlias::parse("devnet-operator_01").is_ok());
    }

    #[test]
    fn ephemeral_wallet_signs_without_exposing_secret_bytes() {
        let alias = crate::WalletAlias::parse("ephemeral")
            .unwrap_or_else(|error| panic!("unexpected alias error: {error}"));
        let wallet = crate::TemporaryWallet::generate(alias);
        let signature = wallet
            .sign_message(b"khadhroony-wallet-test")
            .unwrap_or_else(|error| panic!("unexpected signing error: {error}"));
        assert!(!signature.is_empty());
        assert_eq!(wallet.public_key(), wallet.as_signer().pubkey().to_string());
        assert_eq!(wallet.public_key(), wallet.as_sync_signer().pubkey().to_string());
        assert_eq!(wallet.summary().persistence, crate::WalletPersistence::Temporary);
        assert!(!format!("{wallet:?}").contains("secret"));
    }

    #[tokio::test]
    async fn persistent_wallet_roundtrips_and_preserves_public_key() {
        let directory = tempfile::tempdir()
            .unwrap_or_else(|error| panic!("temporary directory must exist: {error}"));
        let store = crate::TemporaryWalletStore::new(directory.path().join("wallets"))
            .unwrap_or_else(|error| panic!("unexpected store error: {error}"));
        let alias = crate::WalletAlias::parse("integration-wallet")
            .unwrap_or_else(|error| panic!("unexpected alias error: {error}"));
        let created = store
            .create(alias.clone())
            .await
            .unwrap_or_else(|error| panic!("unexpected create error: {error}"));
        let loaded = store
            .load(alias.clone())
            .await
            .unwrap_or_else(|error| panic!("unexpected load error: {error}"));
        assert_eq!(created.public_key(), loaded.public_key());
        assert_eq!(created.summary().persistence, crate::WalletPersistence::Persistent);
        let encoded = std::fs::read(store.wallet_path(&alias))
            .unwrap_or_else(|error| panic!("wallet file must be readable: {error}"));
        let keypair_bytes = serde_json::from_slice::<std::vec::Vec<u8>>(encoded.as_slice())
            .unwrap_or_else(|error| panic!("wallet file must be standard JSON: {error}"));
        assert_eq!(keypair_bytes.len(), crate::SOLANA_KEYPAIR_LENGTH);
    }

    #[tokio::test]
    async fn persistent_wallet_is_never_overwritten() {
        let directory = tempfile::tempdir()
            .unwrap_or_else(|error| panic!("temporary directory must exist: {error}"));
        let store = crate::TemporaryWalletStore::new(directory.path())
            .unwrap_or_else(|error| panic!("unexpected store error: {error}"));
        let alias = crate::WalletAlias::parse("existing")
            .unwrap_or_else(|error| panic!("unexpected alias error: {error}"));
        let first = store
            .create(alias.clone())
            .await
            .unwrap_or_else(|error| panic!("unexpected create error: {error}"));
        let second = store.create(alias.clone()).await;
        assert!(second.is_err());
        let loaded = store
            .load_or_create(alias)
            .await
            .unwrap_or_else(|error| panic!("unexpected load-or-create error: {error}"));
        assert_eq!(first.public_key(), loaded.public_key());
    }

    #[tokio::test]
    async fn corrupted_wallet_file_is_rejected() {
        let directory = tempfile::tempdir()
            .unwrap_or_else(|error| panic!("temporary directory must exist: {error}"));
        let store = crate::TemporaryWalletStore::new(directory.path())
            .unwrap_or_else(|error| panic!("unexpected store error: {error}"));
        let alias = crate::WalletAlias::parse("corrupted")
            .unwrap_or_else(|error| panic!("unexpected alias error: {error}"));
        store
            .create(alias.clone())
            .await
            .unwrap_or_else(|error| panic!("fixture wallet must be created: {error}"));
        std::fs::write(store.wallet_path(&alias), b"[1,2,3]")
            .unwrap_or_else(|error| panic!("fixture must be writable: {error}"));
        let result = store.load(alias).await;
        assert!(result.is_err());
        assert_eq!(
            result
                .err()
                .unwrap_or_else(|| panic!("corrupted wallet must return an error"))
                .code(),
            "wallet_keypair_length_invalid"
        );
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn persistent_wallet_uses_private_unix_permissions() {
        use std::os::unix::fs::PermissionsExt; // rust-rules: trait-import
        let directory = tempfile::tempdir()
            .unwrap_or_else(|error| panic!("temporary directory must exist: {error}"));
        let store = crate::TemporaryWalletStore::new(directory.path().join("private-wallets"))
            .unwrap_or_else(|error| panic!("unexpected store error: {error}"));
        let alias = crate::WalletAlias::parse("permissions")
            .unwrap_or_else(|error| panic!("unexpected alias error: {error}"));
        store
            .create(alias.clone())
            .await
            .unwrap_or_else(|error| panic!("unexpected create error: {error}"));
        let directory_mode = std::fs::metadata(store.directory())
            .unwrap_or_else(|error| panic!("directory metadata must exist: {error}"))
            .permissions()
            .mode()
            & 0o777;
        let file_mode = std::fs::metadata(store.wallet_path(&alias))
            .unwrap_or_else(|error| panic!("file metadata must exist: {error}"))
            .permissions()
            .mode()
            & 0o777;
        assert_eq!(directory_mode, 0o700);
        assert_eq!(file_mode, 0o600);
    }
}