1use anyhow::{anyhow, Context, Result};
2use base64::Engine as _;
3use ed25519_dalek::{Signer, SigningKey};
4
5pub const ENV_KEY_ID: &str = "POLYMARKET_US_KEY_ID";
6pub const ENV_SECRET_KEY: &str = "POLYMARKET_US_SECRET_KEY";
7
8pub const HEADER_ACCESS_KEY: &str = "X-PM-Access-Key";
9pub const HEADER_TIMESTAMP: &str = "X-PM-Timestamp";
10pub const HEADER_SIGNATURE: &str = "X-PM-Signature";
11
12#[derive(Clone)]
13pub struct UsAuth {
14 key_id: String,
15 signing_key: SigningKey,
16}
17
18impl UsAuth {
19 pub fn from_env() -> Result<Self> {
20 let key_id = std::env::var(ENV_KEY_ID).with_context(|| format!("{ENV_KEY_ID} not set"))?;
21 let secret_b64 =
22 std::env::var(ENV_SECRET_KEY).with_context(|| format!("{ENV_SECRET_KEY} not set"))?;
23 Self::from_parts(key_id, &secret_b64)
24 }
25
26 pub fn from_parts(key_id: String, secret_b64: &str) -> Result<Self> {
27 let secret = base64::engine::general_purpose::STANDARD
28 .decode(secret_b64.trim())
29 .context("POLYMARKET_US_SECRET_KEY is not valid Base64")?;
30
31 let signing_key = match secret.len() {
32 64 => {
33 let seed: [u8; 32] = secret[..32].try_into().expect("first 32 bytes");
34 SigningKey::from_bytes(&seed)
35 }
36 32 => {
37 let seed: [u8; 32] = secret.as_slice().try_into().expect("len checked == 32");
38 SigningKey::from_bytes(&seed)
39 }
40 n => {
41 return Err(anyhow!(
42 "POLYMARKET_US_SECRET_KEY must decode to 64 bytes (keypair) or 32 bytes (seed), got {n}"
43 ))
44 }
45 };
46
47 Ok(Self {
48 key_id,
49 signing_key,
50 })
51 }
52
53 pub fn key_id(&self) -> &str {
54 &self.key_id
55 }
56
57 fn signing_payload(timestamp_ms: i64, method: &str, path: &str) -> String {
58 format!("{}{}{}", timestamp_ms, method.to_uppercase(), path)
59 }
60
61 pub fn sign(&self, method: &str, path: &str) -> (i64, String) {
62 let ts = chrono::Utc::now().timestamp_millis();
63 let payload = Self::signing_payload(ts, method, path);
64 let sig_bytes = self.signing_key.sign(payload.as_bytes()).to_bytes();
65 let signature = base64::engine::general_purpose::STANDARD.encode(sig_bytes);
66 (ts, signature)
67 }
68
69 pub fn signed_headers(&self, method: &str, path: &str) -> [(&'static str, String); 3] {
70 let (ts, sig) = self.sign(method, path);
71 [
72 (HEADER_ACCESS_KEY, self.key_id.clone()),
73 (HEADER_TIMESTAMP, ts.to_string()),
74 (HEADER_SIGNATURE, sig),
75 ]
76 }
77}
78
79#[cfg(test)]
80mod tests {
81 use super::*;
82 use ed25519_dalek::{Verifier, VerifyingKey};
83
84 const SAMPLE_SECRET: &str =
85 "lxcsopNhvp+FyZMtVPnHPeHAGihFMPEZcUg6TrJX6kCfwSEXu8v8vmyi3wJbMFUs3a9Fe7mkyRIwfZZkd/5kPg==";
86
87 #[test]
88 fn loads_64_byte_keypair_and_signs_verifiably() {
89 let auth = UsAuth::from_parts("483074f3-key".into(), SAMPLE_SECRET).unwrap();
90
91 let (ts, sig_b64) = auth.sign("GET", "/v1/account/balance");
92 let sig_bytes = base64::engine::general_purpose::STANDARD
93 .decode(sig_b64)
94 .unwrap();
95 let sig = ed25519_dalek::Signature::from_slice(&sig_bytes).unwrap();
96
97 let raw = base64::engine::general_purpose::STANDARD
98 .decode(SAMPLE_SECRET)
99 .unwrap();
100 let pub_bytes: [u8; 32] = raw[32..64].try_into().unwrap();
101 let vk = VerifyingKey::from_bytes(&pub_bytes).unwrap();
102 let payload = format!("{}GET/v1/account/balance", ts);
103 assert!(vk.verify(payload.as_bytes(), &sig).is_ok());
104 }
105}