1use crate::errors::{Error, Result};
4use ed25519_dalek::{Signer, SigningKey};
5use sha2::{Digest, Sha256};
6use std::collections::{BTreeMap, HashMap};
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::sync::{Arc, Mutex, OnceLock};
9use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
10
11pub const API_KEY_ID_ENV: &str = "POLYESTER_API_KEY_ID";
12pub const API_PRIVATE_KEY_ENV: &str = "POLYESTER_API_PRIVATE_KEY";
13pub const ACCOUNT_ID_ENV: &str = "POLYESTER_ACCOUNT_ID";
14
15pub const HEADER_KEY_ID: &str = "X-API-KEY-ID";
16pub const HEADER_TIMESTAMP: &str = "X-API-TIMESTAMP";
17pub const HEADER_SIGNATURE: &str = "X-API-SIGNATURE";
18
19pub const MAX_SIGNING_FUTURE_SKEW_MS: u64 = 5_000;
23const MAX_SIGNING_BACKPRESSURE: Duration = Duration::from_secs(10);
24
25#[derive(Debug, Default)]
26struct SigningTimestampAllocator {
27 last_timestamp_ms: AtomicU64,
28 async_gate: tokio::sync::Mutex<()>,
29}
30
31impl SigningTimestampAllocator {
32 fn next(&self) -> Result<u64> {
33 loop {
34 let now = timestamp_ms_from(SystemTime::now())?;
35 let ceiling = now
36 .checked_add(MAX_SIGNING_FUTURE_SKEW_MS)
37 .ok_or_else(|| Error::transport("signing timestamp ceiling overflow"))?;
38 let observed = self.last_timestamp_ms.load(Ordering::Acquire);
39 let candidate = if observed < now {
40 now
41 } else {
42 observed
43 .checked_add(1)
44 .ok_or_else(|| Error::transport("signing timestamp sequence exhausted"))?
45 };
46
47 if candidate <= ceiling {
48 if self
49 .last_timestamp_ms
50 .compare_exchange_weak(observed, candidate, Ordering::AcqRel, Ordering::Acquire)
51 .is_ok()
52 {
53 return Ok(candidate);
54 }
55 continue;
56 }
57
58 return Err(signing_capacity_error(candidate - ceiling));
59 }
60 }
61
62 async fn next_async(&self) -> Result<u64> {
63 let started = Instant::now();
66 let _gate = tokio::time::timeout(MAX_SIGNING_BACKPRESSURE, self.async_gate.lock())
67 .await
68 .map_err(|_| signing_capacity_error(1))?;
69 loop {
70 match self.next() {
71 Ok(timestamp) => return Ok(timestamp),
72 Err(Error::RateLimit { retry_after, .. })
73 if started.elapsed() < MAX_SIGNING_BACKPRESSURE =>
74 {
75 let wait = Duration::from_secs_f64(retry_after.unwrap_or(0.001).max(0.001));
76 let remaining = MAX_SIGNING_BACKPRESSURE.saturating_sub(started.elapsed());
77 tokio::time::sleep(wait.min(remaining)).await;
78 }
79 Err(error) => return Err(error),
80 }
81 }
82 }
83}
84
85fn signing_capacity_error(wait_ms: u64) -> Error {
86 Error::RateLimit {
87 message: "signing timestamp capacity exhausted; retry after clock advances".to_owned(),
88 retry_after: Some(wait_ms.max(1) as f64 / 1_000.0),
89 }
90}
91
92fn allocator_for_key(key_id: &str) -> Arc<SigningTimestampAllocator> {
93 static ALLOCATORS: OnceLock<Mutex<HashMap<String, Arc<SigningTimestampAllocator>>>> =
94 OnceLock::new();
95 let mut allocators = ALLOCATORS
96 .get_or_init(|| Mutex::new(HashMap::new()))
97 .lock()
98 .unwrap_or_else(|poisoned| poisoned.into_inner());
99 if let Some(existing) = allocators.get(key_id) {
100 return existing.clone();
101 }
102 let allocator = Arc::new(SigningTimestampAllocator::default());
103 allocators.insert(key_id.to_owned(), allocator.clone());
104 allocator
105}
106
107#[derive(Clone)]
113pub struct Credentials {
114 pub key_id: String,
115 signing_key: SigningKey,
116 timestamp_allocator: Arc<SigningTimestampAllocator>,
117}
118
119impl std::fmt::Debug for Credentials {
120 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
121 f.debug_struct("Credentials")
122 .field("key_id", &self.key_id)
123 .field("signing_key", &"<redacted>")
124 .finish()
125 }
126}
127
128impl Credentials {
129 pub fn new(key_id: impl Into<String>, private_key_hex: &str) -> Result<Self> {
130 let key_id = key_id.into().trim().to_owned();
131 if key_id.is_empty() {
132 return Err(Error::auth("API key ID must not be empty"));
133 }
134 let seed = normalize_private_key(private_key_hex)?;
135 let signing_key = SigningKey::from_bytes(&seed);
136 let timestamp_allocator = allocator_for_key(&key_id);
137 Ok(Self {
138 key_id,
139 signing_key,
140 timestamp_allocator,
141 })
142 }
143
144 pub fn load(
146 api_key_id: Option<&str>,
147 api_private_key: Option<&str>,
148 from_env: bool,
149 ) -> Result<Option<Self>> {
150 let mut key_id = api_key_id.unwrap_or("").trim().to_owned();
151 let mut private = api_private_key.unwrap_or("").trim().to_owned();
152 if from_env {
153 if key_id.is_empty() {
154 key_id = std::env::var(API_KEY_ID_ENV)
155 .unwrap_or_default()
156 .trim()
157 .to_owned();
158 }
159 if private.is_empty() {
160 private = std::env::var(API_PRIVATE_KEY_ENV)
161 .unwrap_or_default()
162 .trim()
163 .to_owned();
164 }
165 }
166 if key_id.is_empty() && private.is_empty() {
167 return Ok(None);
168 }
169 if key_id.is_empty() || private.is_empty() {
170 let msg = if from_env {
171 "Both POLYESTER_API_KEY_ID and POLYESTER_API_PRIVATE_KEY are required"
172 } else {
173 "Both api_key_id and api_private_key are required"
174 };
175 return Err(Error::auth(msg));
176 }
177 Ok(Some(Self::new(key_id, &private)?))
178 }
179
180 pub fn sign_request(
181 &self,
182 method: &str,
183 raw_url: &str,
184 body: &[u8],
185 timestamp_ms: Option<&str>,
186 ) -> Result<BTreeMap<String, String>> {
187 let ts = match timestamp_ms {
188 Some(value) => value.to_owned(),
189 None => self.timestamp_allocator.next()?.to_string(),
190 };
191 let canonical = canonical_signing_string(&ts, method, raw_url, body)?;
192 let sig = self.signing_key.sign(canonical.as_bytes());
193 let mut headers = BTreeMap::new();
194 headers.insert(HEADER_KEY_ID.to_owned(), self.key_id.clone());
195 headers.insert(HEADER_TIMESTAMP.to_owned(), ts);
196 headers.insert(HEADER_SIGNATURE.to_owned(), hex::encode(sig.to_bytes()));
197 Ok(headers)
198 }
199
200 pub fn sign_payload(&self, payload: &[u8]) -> Vec<u8> {
206 self.signing_key.sign(payload).to_bytes().to_vec()
207 }
208
209 pub async fn sign_request_async(
215 &self,
216 method: &str,
217 raw_url: &str,
218 body: &[u8],
219 timestamp_ms: Option<&str>,
220 ) -> Result<BTreeMap<String, String>> {
221 let ts = match timestamp_ms {
222 Some(value) => value.to_owned(),
223 None => self.timestamp_allocator.next_async().await?.to_string(),
224 };
225 let canonical = canonical_signing_string(&ts, method, raw_url, body)?;
226 let sig = self.signing_key.sign(canonical.as_bytes());
227 let mut headers = BTreeMap::new();
228 headers.insert(HEADER_KEY_ID.to_owned(), self.key_id.clone());
229 headers.insert(HEADER_TIMESTAMP.to_owned(), ts);
230 headers.insert(HEADER_SIGNATURE.to_owned(), hex::encode(sig.to_bytes()));
231 Ok(headers)
232 }
233}
234
235fn timestamp_ms_from(now: SystemTime) -> Result<u64> {
236 let elapsed = now
237 .duration_since(UNIX_EPOCH)
238 .map_err(|_| Error::transport("system clock is before UNIX_EPOCH"))?;
239 u64::try_from(elapsed.as_millis())
240 .map_err(|_| Error::transport("Unix timestamp milliseconds exceed u64 range"))
241}
242
243pub fn normalize_private_key(value: &str) -> Result<[u8; 32]> {
245 let private = hex::decode(value.trim())
246 .map_err(|_| Error::auth("API private key must be a valid hex string or raw bytes"))?;
247 if private.len() != 32 {
248 return Err(Error::auth(
249 "Ed25519 API private key must be exactly 32 bytes",
250 ));
251 }
252 let mut seed = [0u8; 32];
253 seed.copy_from_slice(&private);
254 Ok(seed)
255}
256
257pub fn account_id_from_env() -> Option<String> {
258 let v = std::env::var(ACCOUNT_ID_ENV).ok()?.trim().to_owned();
259 if v.is_empty() { None } else { Some(v) }
260}
261
262const QUERY_COMPONENT_ASCII_SET: &percent_encoding::AsciiSet = &percent_encoding::NON_ALPHANUMERIC
269 .remove(b'-')
270 .remove(b'_')
271 .remove(b'.')
272 .remove(b'~');
273
274pub fn encode_query_component(s: &str) -> String {
276 percent_encoding::utf8_percent_encode(s, QUERY_COMPONENT_ASCII_SET).to_string()
277}
278
279fn canonical_query_from_url(parsed: &url::Url) -> String {
281 let mut pairs: Vec<(String, String)> = parsed
282 .query_pairs()
283 .map(|(k, v)| (k.into_owned(), v.into_owned()))
284 .collect();
285 pairs.sort_by(|a, b| a.0.cmp(&b.0).then(a.1.cmp(&b.1)));
286 pairs
287 .into_iter()
288 .map(|(k, v)| {
289 format!(
290 "{}={}",
291 encode_query_component(&k),
292 encode_query_component(&v)
293 )
294 })
295 .collect::<Vec<_>>()
296 .join("&")
297}
298
299pub fn canonical_query(raw_url: &str) -> Result<String> {
300 let parsed = url::Url::parse(raw_url)
301 .map_err(|err| Error::validation(format!("invalid signing URL: {err}")))?;
302 Ok(canonical_query_from_url(&parsed))
303}
304
305pub fn canonical_signing_string(
306 timestamp_ms: &str,
307 method: &str,
308 raw_url: &str,
309 body: &[u8],
310) -> Result<String> {
311 let parsed = url::Url::parse(raw_url)
312 .map_err(|err| Error::validation(format!("invalid signing URL: {err}")))?;
313 let pathname = {
314 let path = parsed.path();
315 if path.is_empty() {
316 "/".to_owned()
317 } else {
318 path.to_owned()
319 }
320 };
321 let sum = Sha256::digest(body);
322 Ok([
323 timestamp_ms,
324 &method.to_uppercase(),
325 &pathname,
326 &canonical_query_from_url(&parsed),
327 &hex::encode(sum),
328 ]
329 .join("\n"))
330}
331
332pub fn request_url(api_base: &str, procedure: &str) -> String {
333 let base = api_base.trim_end_matches('/');
334 let proc = if procedure.starts_with('/') {
335 procedure.to_owned()
336 } else {
337 format!("/{procedure}")
338 };
339 format!("{base}{proc}")
340}
341
342pub fn generate_ed25519_keypair() -> (String, String) {
344 use ed25519_dalek::SigningKey;
345 use rand_core::OsRng;
346 let signing = SigningKey::generate(&mut OsRng);
347 let seed = signing.to_bytes();
348 let public = signing.verifying_key().to_bytes();
349 (hex::encode(seed), hex::encode(public))
350}
351
352#[cfg(test)]
353mod tests {
354 use super::*;
355
356 #[test]
357 fn canonical_query_sorts_and_encodes_values() {
358 let got = canonical_query("https://api.example.test/path?b=2&a=hello world").unwrap();
359 assert_eq!(got, "a=hello%20world&b=2");
360 }
361
362 #[test]
363 fn encode_query_component_preserves_rfc3986_unreserved() {
364 assert_eq!(encode_query_component("api-keys"), "api-keys");
365 assert_eq!(encode_query_component("a_b.c~d-e"), "a_b.c~d-e");
366 assert_eq!(encode_query_component("hello world"), "hello%20world");
367 assert_eq!(encode_query_component("a+b"), "a%2Bb");
368 assert_eq!(
369 encode_query_component("private:auth:api-keys:account:proto"),
370 "private%3Aauth%3Aapi-keys%3Aaccount%3Aproto"
371 );
372 }
373
374 #[test]
375 fn canonical_query_preserves_hyphens_in_channel_param() {
376 let url =
377 "https://api.example.test/v1/rt/subscribe?channel=private:auth:api-keys:account:proto";
378 assert_eq!(
379 canonical_query(url).unwrap(),
380 "channel=private%3Aauth%3Aapi-keys%3Aaccount%3Aproto"
381 );
382 }
383
384 #[test]
385 fn canonical_query_shared_vectors() {
386 let cases = [
389 (
390 "https://api.example.test/x?z=1&a=hello world&m=a+b",
391 "a=hello%20world&m=a%20b&z=1",
392 ),
393 (
394 "https://api.example.test/x?z=1&a=hello%20world&m=a%2Bb",
395 "a=hello%20world&m=a%2Bb&z=1",
396 ),
397 ("https://api.example.test/x?b=&a=1", "a=1&b="),
398 ("https://api.example.test/x?a=1&a=2&b=0", "a=1&a=2&b=0"),
399 (
400 "https://api.example.test/x?path=foo/bar&name=a_b.c~d-e",
401 "name=a_b.c~d-e&path=foo%2Fbar",
402 ),
403 (
404 "https://api.example.test/x?msg=%E2%9C%93&plain=ok",
405 "msg=%E2%9C%93&plain=ok",
406 ),
407 ];
408 for (url, want) in cases {
409 assert_eq!(canonical_query(url).unwrap(), want, "url={url}");
410 }
411 }
412
413 #[test]
414 fn canonical_signing_string_matches_contract() {
415 let got = canonical_signing_string(
416 "123",
417 "post",
418 "https://api.example.test/foo/bar?b=2&a=1",
419 b"{}",
420 )
421 .unwrap();
422 let want = "123\nPOST\n/foo/bar\na=1&b=2\n44136fa355b3678a1146ad16f7e8649e94fb4fc21fe77e8310c060f61caaff8a";
423 assert_eq!(got, want);
424 }
425
426 #[test]
427 fn canonical_string_empty_body() {
428 let s = canonical_signing_string(
429 "1700000000000",
430 "POST",
431 "https://api-devnet.polyester.ai/orders.v1.OrdersService/CreateOrder",
432 b"",
433 )
434 .unwrap();
435 let expected_hash = hex::encode(Sha256::digest(b""));
436 assert!(s.contains(&expected_hash));
437 assert!(s.starts_with("1700000000000\nPOST\n/orders.v1.OrdersService/CreateOrder\n\n"));
438 }
439
440 #[test]
441 fn sign_request_returns_polyester_headers() {
442 let (seed, _) = generate_ed25519_keypair();
443 let creds = Credentials::new("key_123", &seed).unwrap();
444 let headers = creds
445 .sign_request("POST", "https://api.example.test/foo", b"{}", Some("123"))
446 .unwrap();
447 assert_eq!(headers.get(HEADER_KEY_ID).unwrap(), "key_123");
448 assert_eq!(headers.get(HEADER_TIMESTAMP).unwrap(), "123");
449 assert_eq!(headers.get(HEADER_SIGNATURE).unwrap().len(), 128);
450 }
451
452 #[tokio::test(flavor = "current_thread")]
453 async fn ten_thousand_identical_requests_get_unique_bounded_auth_tuples_without_blocking_runtime()
454 {
455 use std::collections::HashSet;
456
457 let (seed, _) = generate_ed25519_keypair();
458 let creds = Credentials::new("key_123", &seed).unwrap();
459 let before = timestamp_ms_from(SystemTime::now()).unwrap();
460 let ticker = tokio::spawn(async {
461 let mut previous = tokio::time::Instant::now();
462 let mut largest_gap = Duration::ZERO;
463 for _ in 0..100 {
464 tokio::time::sleep(Duration::from_millis(10)).await;
465 let now = tokio::time::Instant::now();
466 largest_gap = largest_gap.max(now.duration_since(previous));
467 previous = now;
468 }
469 largest_gap
470 });
471 let mut headers = Vec::with_capacity(10_000);
475 for chunk_start in (0..10_000).step_by(250) {
476 let chunk_end = (chunk_start + 250).min(10_000);
477 let handles = (chunk_start..chunk_end)
478 .map(|_| {
479 let creds = creds.clone();
480 tokio::spawn(async move {
481 let headers = creds
482 .sign_request_async("POST", "https://api.example.test/foo", b"{}", None)
483 .await
484 .unwrap();
485 let observed_at_ms = timestamp_ms_from(SystemTime::now()).unwrap();
486 (headers, observed_at_ms)
487 })
488 })
489 .collect::<Vec<_>>();
490 for handle in handles {
491 headers.push(handle.await.unwrap());
492 }
493 tokio::task::yield_now().await;
494 }
495 let largest_timer_gap = ticker.await.unwrap();
496 assert_eq!(headers.len(), 10_000);
497 assert!(
498 largest_timer_gap < Duration::from_secs(1),
499 "signing backpressure stalled the current-thread runtime for {largest_timer_gap:?}"
500 );
501 let mut timestamps = HashSet::with_capacity(headers.len());
502 let mut signatures = HashSet::with_capacity(headers.len());
503 for (item, observed_at_ms) in headers {
504 let timestamp = item[HEADER_TIMESTAMP].parse::<u64>().unwrap();
505 assert!(timestamp >= before);
506 assert!(timestamp <= observed_at_ms + MAX_SIGNING_FUTURE_SKEW_MS);
507 assert!(
508 timestamps.insert(timestamp),
509 "duplicate timestamp {timestamp}"
510 );
511 assert!(
512 signatures.insert(item[HEADER_SIGNATURE].clone()),
513 "duplicate signature for timestamp {timestamp}"
514 );
515 }
516 assert_eq!(timestamps.len(), 10_000);
517 assert_eq!(signatures.len(), 10_000);
518 }
519
520 #[tokio::test]
521 async fn independently_constructed_credentials_share_process_timestamp_allocator() {
522 let (seed, _) = generate_ed25519_keypair();
523 let first = Credentials::new("shared-key", &seed).unwrap();
524 let second = Credentials::new("shared-key", &seed).unwrap();
525 assert!(Arc::ptr_eq(
526 &first.timestamp_allocator,
527 &second.timestamp_allocator
528 ));
529 let (first, second) = tokio::join!(
530 first.sign_request_async("POST", "https://api.test/order", b"{}", None),
531 second.sign_request_async("POST", "https://api.test/order", b"{}", None)
532 );
533 assert_ne!(
534 first.unwrap()[HEADER_TIMESTAMP],
535 second.unwrap()[HEADER_TIMESTAMP]
536 );
537
538 let transient = Credentials::new("process-lifetime-key", &seed).unwrap();
539 let before_drop = transient
540 .sign_request("POST", "https://api.test/order", b"{}", None)
541 .unwrap()[HEADER_TIMESTAMP]
542 .parse::<u64>()
543 .unwrap();
544 drop(transient);
545 let reconstructed = Credentials::new("process-lifetime-key", &seed).unwrap();
546 let after_drop = reconstructed
547 .sign_request("POST", "https://api.test/order", b"{}", None)
548 .unwrap()[HEADER_TIMESTAMP]
549 .parse::<u64>()
550 .unwrap();
551 assert!(after_drop > before_drop);
552 }
553
554 #[test]
555 fn synchronous_signing_capacity_returns_retryable_error_without_sleeping() {
556 let allocator = SigningTimestampAllocator::default();
557 allocator
560 .last_timestamp_ms
561 .store(u64::MAX - 2, Ordering::Release);
562 let started = Instant::now();
563 let error = allocator.next().unwrap_err();
564 assert!(matches!(error, Error::RateLimit { .. }));
565 assert!(error.is_retryable());
566 assert!(error.retry_after().is_some_and(|value| value > 0.0));
567 assert!(
568 started.elapsed() < Duration::from_millis(50),
569 "synchronous capacity handling blocked for {:?}",
570 started.elapsed()
571 );
572 }
573
574 #[test]
575 fn round_trip_credentials() {
576 let (seed, _) = generate_ed25519_keypair();
577 let creds = Credentials::new("ak_test", &seed).unwrap();
578 let headers = creds
579 .sign_request(
580 "POST",
581 "https://api-devnet.polyester.ai/orders.v1.OrdersService/CreateOrder",
582 b"{}",
583 Some("1"),
584 )
585 .unwrap();
586 assert_eq!(headers.get(HEADER_KEY_ID).unwrap(), "ak_test");
587 assert_eq!(headers.get(HEADER_TIMESTAMP).unwrap(), "1");
588 assert_eq!(headers.get(HEADER_SIGNATURE).unwrap().len(), 128);
589 }
590
591 #[test]
592 fn credentials_reject_empty_key_id() {
593 let (seed, _) = generate_ed25519_keypair();
594 let err = Credentials::new(" ", &seed).unwrap_err();
595 assert!(matches!(err, Error::Auth(_)));
596 }
597
598 #[test]
599 fn load_credentials_requires_both() {
600 let err = Credentials::load(Some("ak_test"), Some(""), false).unwrap_err();
601 assert!(matches!(err, Error::Auth(_)));
602 }
603
604 #[test]
605 fn load_credentials_none_when_empty() {
606 assert!(Credentials::load(None, None, false).unwrap().is_none());
607 }
608
609 #[test]
610 fn request_url_joins_base_and_procedure() {
611 assert_eq!(
612 request_url(
613 "https://api.example.test/",
614 "orders.v1.OrdersService/CreateOrder"
615 ),
616 "https://api.example.test/orders.v1.OrdersService/CreateOrder"
617 );
618 assert_eq!(
619 request_url("https://api.example.test", "/auth.v1.AuthService/Me"),
620 "https://api.example.test/auth.v1.AuthService/Me"
621 );
622 }
623
624 #[test]
625 fn signing_rejects_unparseable_urls() {
626 let (seed, _) = generate_ed25519_keypair();
627 let creds = Credentials::new("ak_test", &seed).unwrap();
628 let err = creds
629 .sign_request("POST", "not a url", b"{}", Some("1"))
630 .unwrap_err();
631 assert!(matches!(err, Error::Validation(_)));
632 assert!(canonical_query("not a url").is_err());
633 assert!(canonical_signing_string("1", "POST", "not a url", b"").is_err());
634 }
635
636 #[test]
637 fn pre_epoch_clock_is_an_error_not_a_panic() {
638 let before_epoch = UNIX_EPOCH
639 .checked_sub(std::time::Duration::from_secs(1))
640 .unwrap();
641 assert!(matches!(
642 timestamp_ms_from(before_epoch),
643 Err(Error::Transport(_))
644 ));
645 }
646}