Skip to main content

polymarket_us/
auth.rs

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}