aimo_core/
keys.rs

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

aimo_core/
keys.rs
