aimo_core/
keys.rs

1use std::str::FromStr;
2
3use anyhow::{anyhow, bail, Ok};
4use chrono::{DateTime, Utc};
5use schemars::JsonSchema;
6use serde::{Deserialize, Serialize};
7use sha2::{Digest, Sha256};
8use solana_sdk::{pubkey::Pubkey, signature::Signature};
9
10/// Secret key format
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct SecretKeyV1 {
13    pub version: u8,          // 1 byte
14    pub wallet: Wallet,       // use enum: 1 byte
15    pub signer: String,       // 32 bytes
16    pub signature: String,    // 64 bytes
17    pub metadata: MetadataV1, // 28 bytes
18}
19
20impl SecretKeyV1 {
21    /// Encode the secret key into a string in the form of:
22    ///
23    /// `aimo-sk-{scope}-{base58_encoded_secret_key_json}`
24    pub fn into_string(self, scope: &str) -> anyhow::Result<String> {
25        let raw = SecretKeyRawV1::try_from(self)?;
26        let base58_encoded = bs58::encode(raw.into_bytes())
27            // .with_check()
28            .into_string();
29
30        Ok(format!("aimo-sk-{scope}-{base58_encoded}"))
31    }
32
33    fn split_sk_string(sk: &str) -> Option<(&str, &str)> {
34        let mut parts = sk.splitn(4, '-');
35        let aimo = parts.next()?;
36        if aimo != "aimo" {
37            return None;
38        }
39        let prefix = parts.next()?;
40        if prefix != "sk" {
41            return None;
42        }
43
44        let scope = parts.next()?;
45        let base58_value = parts.next()?;
46
47        Some((scope, base58_value))
48    }
49
50    pub fn decode(sk: &str) -> anyhow::Result<(String, Self)> {
51        let (scope, key) = Self::split_sk_string(sk).ok_or(anyhow!(
52            "Invalid secret key: Failed to split secret key into valid parts"
53        ))?;
54        let decoded_bytes = bs58::decode(key).into_vec()?;
55        let raw = SecretKeyRawV1::from_bytes(&decoded_bytes[..])?;
56
57        Ok((scope.to_string(), raw.try_into()?))
58    }
59
60    pub fn verify_signature(&self) -> anyhow::Result<()> {
61        // Create canonical JSON string of the metadata
62        let canonical_metadata = self.metadata.to_canonical_json()?;
63        let public_key = Pubkey::from_str(&self.signer)?;
64        let signature = Signature::from_str(&self.signature)?;
65        let is_valid = signature.verify(public_key.as_ref(), canonical_metadata.as_bytes());
66
67        if !is_valid {
68            bail!("Wrong signature");
69        }
70
71        // Check expiry
72        if let Some(dt) = DateTime::<Utc>::from_timestamp_millis(
73            self.metadata.created_at + self.metadata.valid_for,
74        ) {
75            if dt < Utc::now() {
76                bail!("Expired");
77            }
78        } else {
79            bail!("Invalid timestamp");
80        }
81
82        Ok(())
83    }
84
85    pub fn into_hash(self) -> anyhow::Result<[u8; 32]> {
86        let bytes = SecretKeyRawV1::try_from(self)?.into_bytes();
87        let mut hasher = Sha256::new();
88        hasher.update(&bytes[..]);
89        Ok(hasher.finalize().into())
90    }
91}
92
93#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
94pub struct MetadataV1 {
95    pub created_at: i64,
96    pub valid_for: i64,
97    pub usage_limit: u64,
98    pub scopes: Vec<Scope>,
99}
100
101impl MetadataV1 {
102    /// Encode the metadata to a canonical JSON string for signing
103    pub fn to_canonical_json(&self) -> anyhow::Result<String> {
104        let metadata_json = serde_json::to_value(self)?;
105        canonical_json::to_string(&metadata_json).map_err(Into::into)
106    }
107}
108
109#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
110pub enum Wallet {
111    #[serde(rename = "solana")]
112    Solana,
113}
114
115#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
116pub enum Scope {
117    #[serde(rename = "completion_model")]
118    CompletionModel,
119}
120
121impl FromStr for Scope {
122    type Err = anyhow::Error;
123
124    fn from_str(s: &str) -> Result<Self, Self::Err> {
125        match s {
126            "completion_model" => Ok(Self::CompletionModel),
127            _ => Err(anyhow!("Scope {s} not supported")),
128        }
129    }
130}
131
132type WalletEnum = u8;
133type ScopeBitMap = u64;
134
135#[derive(Debug, Clone)]
136pub struct SecretKeyRawV1 {
137    pub version: u8,             // 1 byte
138    pub wallet: WalletEnum,      // 1 byte
139    pub signer: [u8; 32],        // 32 bytes
140    pub signature: [u8; 64],     // 64 bytes
141    pub metadata: MetadataRawV1, // 32 bytes
142}
143
144impl SecretKeyRawV1 {
145    pub const BYTES: usize = 130; // 1 + 1 + 32 + 64 + 32
146
147    pub fn into_bytes(self) -> Vec<u8> {
148        [
149            vec![self.version, self.wallet],
150            self.signer.to_vec(),
151            self.signature.to_vec(),
152            self.metadata.into_bytes(),
153        ]
154        .concat()
155    }
156
157    pub fn from_bytes(bytes: &[u8]) -> anyhow::Result<Self> {
158        if bytes.len() != Self::BYTES {
159            bail!(
160                "Bytes length doesn't match: expect {}, got {}",
161                Self::BYTES,
162                bytes.len()
163            );
164        }
165
166        let version = bytes[0];
167        let wallet = bytes[1];
168        let signer: [u8; 32] = bytes[2..34].try_into()?;
169        let signature: [u8; 64] = bytes[34..98].try_into()?;
170        let metadata = MetadataRawV1::from_bytes(&bytes[98..130])?;
171
172        Ok(Self {
173            version,
174            wallet,
175            signer,
176            signature,
177            metadata,
178        })
179    }
180}
181
182impl TryFrom<SecretKeyV1> for SecretKeyRawV1 {
183    type Error = anyhow::Error;
184
185    fn try_from(value: SecretKeyV1) -> Result<Self, Self::Error> {
186        if value.wallet != Wallet::Solana {
187            bail!("Wallet not supported");
188        }
189
190        // Decode with base58 for solana wallets
191        let signer: [u8; 32] = bs58::decode(&value.signer).into_vec()?[..].try_into()?;
192        let signature: [u8; 64] = bs58::decode(&value.signature).into_vec()?[..].try_into()?;
193
194        // TODO: Handle more wallets
195        let wallet: WalletEnum = wallets::SOLANA;
196
197        Ok(Self {
198            version: value.version,
199            wallet,
200            signer,
201            signature,
202            metadata: value.metadata.try_into()?,
203        })
204    }
205}
206
207impl TryFrom<SecretKeyRawV1> for SecretKeyV1 {
208    type Error = anyhow::Error;
209
210    fn try_from(value: SecretKeyRawV1) -> Result<Self, Self::Error> {
211        if value.wallet >= wallets::TOTAL_WALLETS_SUPPORTED {
212            bail!("Unsupported wallet type in secret key: {}", value.wallet);
213        }
214
215        let wallet = match value.wallet {
216            wallets::SOLANA => Wallet::Solana,
217
218            // This will never happen
219            _ => Wallet::Solana,
220        };
221
222        let signer = bs58::encode(&value.signer).into_string();
223        let signature = bs58::encode(&value.signature).into_string();
224
225        Ok(Self {
226            version: value.version,
227            wallet,
228            signer,
229            signature,
230            metadata: value.metadata.try_into()?,
231        })
232    }
233}
234
235#[derive(Debug, Clone, Copy)]
236pub struct MetadataRawV1 {
237    pub created_at: i64,     // 8 bytes
238    pub valid_for: i64,      // 8 bytes
239    pub usage_limit: u64,    // 8 bytes
240    pub scopes: ScopeBitMap, // 8 bytes
241}
242
243impl MetadataRawV1 {
244    pub const BYTES: usize = 32; // 8 + 8 + 8 + 8
245
246    pub fn into_bytes(self) -> Vec<u8> {
247        // We want the bytes to be serialized "from left to right"
248        // So we convert the numbers into big-endian bytes
249
250        [
251            self.created_at.to_be_bytes().to_vec(),
252            self.valid_for.to_be_bytes().to_vec(),
253            self.usage_limit.to_be_bytes().to_vec(),
254            self.scopes.to_be_bytes().to_vec(),
255        ]
256        .concat()
257    }
258
259    pub fn from_bytes(bytes: &[u8]) -> anyhow::Result<Self> {
260        if bytes.len() != Self::BYTES {
261            bail!(
262                "Bytes length doesn't match: expect {}, got {}",
263                Self::BYTES,
264                bytes.len()
265            );
266        }
267
268        let created_at = i64::from_be_bytes(bytes[0..8].try_into()?);
269        let valid_for = i64::from_be_bytes(bytes[8..16].try_into()?);
270        let usage_limit = u64::from_be_bytes(bytes[16..24].try_into()?);
271        let scopes = u64::from_be_bytes(bytes[24..32].try_into()?);
272
273        Ok(Self {
274            created_at,
275            valid_for,
276            usage_limit,
277            scopes,
278        })
279    }
280}
281
282impl TryFrom<MetadataV1> for MetadataRawV1 {
283    type Error = anyhow::Error;
284
285    /// Keep `TryFrom` here even though this doesn't produce errors now
286    fn try_from(value: MetadataV1) -> Result<Self, Self::Error> {
287        // Convert options list into a bitmap
288        let bitmap: ScopeBitMap = value.scopes.iter().fold(0, |bm, scope| match scope {
289            Scope::CompletionModel => bm | 1 << scopes::COMPLETION_MODEL,
290        });
291
292        Ok(Self {
293            created_at: value.created_at,
294            valid_for: value.valid_for,
295            usage_limit: value.usage_limit,
296            scopes: bitmap,
297        })
298    }
299}
300
301impl TryFrom<MetadataRawV1> for MetadataV1 {
302    type Error = anyhow::Error;
303
304    fn try_from(value: MetadataRawV1) -> Result<Self, Self::Error> {
305        // Scopes       Enabled
306        // Supported    | 0     | 1
307        //          0   | 1     | 0
308        //          1   | 1     | 1
309        // ------------------------
310        // > is_valid = supported | !enabled
311        // > is_invalid = !(supported | !enabled)
312        //              = !supported & enabled
313        if (!scopes::SCOPES_SUPPORTED & value.scopes) > 0 {
314            bail!("Secret key contains currently unsupported scope type");
315        }
316
317        let scopes = if value.scopes | scopes::COMPLETION_MODEL > 0 {
318            vec![Scope::CompletionModel]
319        } else {
320            vec![]
321        };
322
323        Ok(Self {
324            created_at: value.created_at,
325            valid_for: value.valid_for,
326            usage_limit: value.usage_limit,
327            scopes,
328        })
329    }
330}
331
332pub mod wallets {
333    //! Signing wallets supported: `0x00` - `0xFF`
334
335    use super::WalletEnum;
336    pub const TOTAL_WALLETS_SUPPORTED: WalletEnum = 1;
337
338    /// The default option: Solana wallets
339    pub const SOLANA: WalletEnum = 0x00;
340}
341
342pub mod scopes {
343    //! Scope bitmap options: 0 - 31, from lower bit to higher bit
344    //!
345
346    use super::ScopeBitMap;
347
348    pub const SCOPES_SUPPORTED: ScopeBitMap = 0x01;
349
350    /// Scope: `model:completion`
351    ///
352    /// Position: `0x01` (1 << 0)
353    pub const COMPLETION_MODEL: ScopeBitMap = 0;
354}
aimo_core/keys.rs

aimo_core/
keys.rs
