solagent_wallet_solana/
lib.rs

1// Copyright 2025 zTgx
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use dotenv::dotenv;
16use solana_sdk::{bs58, pubkey::Pubkey, signature::Keypair, signer::Signer};
17use std::env;
18
19/// Represents a wallet containing a keypair and its corresponding public key.
20#[derive(Debug)]
21pub struct Wallet {
22    /// The keypair associated with the wallet.  This contains the private key.
23    pub keypair: Keypair,
24    /// The public key associated with the wallet.
25    pub pubkey: Pubkey,
26}
27
28impl Default for Wallet {
29    /// Creates a new wallet with a randomly generated keypair.
30    fn default() -> Self {
31        Self::new()
32    }
33}
34
35impl Wallet {
36    /// Creates a new wallet with a randomly generated keypair.
37    pub fn new() -> Self {
38        let keypair = Keypair::new();
39        let pubkey = keypair.pubkey();
40        Self { keypair, pubkey }
41    }
42
43    /// Creates a wallet from a private key stored in an environment variable.
44    ///
45    /// This function reads the environment variable specified by `variable_name`,
46    /// decodes the base58 encoded private key, and creates a `Wallet` instance.
47    ///
48    /// # Arguments
49    ///
50    /// * `variable_name` - The name of the environment variable containing the private key.
51    ///
52    /// # Returns
53    ///
54    /// * `Ok(Wallet)` - If the wallet was successfully created.
55    /// * `Err(String)` - If the environment variable is not found or the private key is invalid.
56    pub fn from_env(variable_name: &str) -> Result<Self, String> {
57        dotenv().ok(); // Load environment variables from .env file (if present)
58
59        let private_key =
60            env::var(variable_name).map_err(|_| format!("Environment variable '{}' not found", variable_name))?;
61
62        Self::from_base58(&private_key)
63    }
64
65    /// Creates a wallet from a base58 encoded private key.
66    ///
67    /// # Arguments
68    ///
69    /// * `private_key` - The base58 encoded private key.
70    ///
71    /// # Returns
72    ///
73    /// * `Ok(Wallet)` - If the wallet was successfully created.
74    /// * `Err(String)` - If the private key is invalid or not properly encoded.
75    pub fn from_base58(private_key: &str) -> Result<Self, String> {
76        let secret_key = bs58::decode(private_key).into_vec().map_err(|_| "Invalid base58 private key".to_string())?;
77
78        let keypair = Keypair::from_bytes(&secret_key).map_err(|_| "Invalid private key bytes".to_string())?;
79
80        let pubkey = keypair.pubkey();
81        Ok(Self { keypair, pubkey })
82    }
83
84    /// Returns the base58 encoded private key of the wallet.
85    pub fn to_base58(&self) -> String {
86        self.keypair.to_base58_string()
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93    use std::env;
94
95    #[test]
96    fn test_wallet_from_env_valid() {
97        // Set an environment variable for the test
98        let keypair = Keypair::new();
99        let private_key = keypair.to_base58_string();
100        env::set_var("TEST_PRIVATE_KEY", &private_key);
101
102        let wallet = Wallet::from_env("TEST_PRIVATE_KEY").unwrap();
103        assert_eq!(wallet.pubkey, keypair.pubkey());
104
105        // Clean up the environment variable after the test
106        env::remove_var("TEST_PRIVATE_KEY");
107    }
108
109    #[test]
110    fn test_wallet_from_env_not_found() {
111        let result = Wallet::from_env("NON_EXISTENT_VARIABLE");
112        assert!(result.is_err());
113        assert_eq!(result.unwrap_err(), "Environment variable 'NON_EXISTENT_VARIABLE' not found");
114    }
115
116    #[test]
117    fn test_wallet_creation() {
118        let wallet = Wallet::new();
119        assert_ne!(wallet.pubkey, Pubkey::default());
120    }
121
122    #[test]
123    fn test_wallet_from_base58_valid() {
124        let original_keypair = Keypair::new();
125        let private_key = original_keypair.to_base58_string();
126        let wallet = Wallet::from_base58(&private_key).unwrap();
127        assert_eq!(wallet.pubkey, original_keypair.pubkey());
128    }
129
130    #[test]
131    fn test_wallet_from_base58_invalid_base58() {
132        let result = Wallet::from_base58("invalid_key");
133        assert!(result.is_err());
134        assert_eq!(result.unwrap_err(), "Invalid base58 private key");
135    }
136
137    #[test]
138    fn test_wallet_from_base58_invalid_bytes() {
139        // Create a base58 string that's the wrong length to be a key
140        let invalid_bytes = bs58::encode([0u8; 10]).into_string(); //Incorrect number of bytes
141        let result = Wallet::from_base58(&invalid_bytes);
142        assert!(result.is_err());
143        assert_eq!(result.unwrap_err(), "Invalid private key bytes");
144    }
145
146    #[test]
147    fn test_wallet_to_base58() {
148        let wallet = Wallet::new();
149        let base58_key = wallet.to_base58();
150        assert!(!base58_key.is_empty());
151
152        let wallet2 = Wallet::from_base58(&base58_key).unwrap();
153        assert_eq!(wallet.pubkey, wallet2.pubkey);
154    }
155}
solagent_wallet_solana/lib.rs

solagent_wallet_solana/
lib.rs
