1use aes::cipher::{KeyIvInit, StreamCipher};
42use base64::{engine::general_purpose::URL_SAFE, Engine as _};
43use byteorder::{BigEndian, ByteOrder};
44use hmac::{Hmac, Mac};
45use sha2::Sha256;
46use std::io::Write;
47use thiserror::Error;
48use time::{Duration, OffsetDateTime};
49
50type HmacSha256 = Hmac<Sha256>;
51type Aes256Ctr64BE = ctr::Ctr64BE<aes::Aes256>;
52
53const UNIX_EPOCH: OffsetDateTime = time::OffsetDateTime::UNIX_EPOCH;
54
55#[derive(Clone, Debug)]
64pub struct Keys {
65 pub encryption_key: [u8; 32],
67 pub integrity_key: [u8; 32],
69}
70
71impl Keys {
72 pub fn new(encryption_key: &[u8], integrity_key: &[u8]) -> Result<Self, CryptoError> {
90 let encryption_key: [u8; 32] = encryption_key
91 .try_into()
92 .map_err(|_| CryptoError::InvalidKey)?;
93 let integrity_key: [u8; 32] = integrity_key
94 .try_into()
95 .map_err(|_| CryptoError::InvalidKey)?;
96
97 Ok(Self {
98 encryption_key,
99 integrity_key,
100 })
101 }
102}
103
104#[derive(Error, Debug)]
106pub enum CryptoError {
107 #[error("invalid key")]
109 InvalidKey,
110 #[error("invalid signature")]
112 InvalidSign,
113 #[error("invalid init vector")]
115 InvalidInitVector,
116 #[error("data too short")]
118 DataTooShort,
119 #[error("payload size mismatch")]
121 PayloadSizeMismatch,
122 #[error("decode error: {0}")]
124 DecodeError(#[from] base64::DecodeError),
125 #[error("io error: {0}")]
127 IoError(#[from] std::io::Error),
128}
129
130pub struct Crypto {
144 pub keys: Keys,
146}
147
148impl Crypto {
149 pub fn new(keys: Keys) -> Self {
160 Self { keys }
161 }
162
163 pub const IV_BASE: usize = 0;
165 pub const IV_SIZE: usize = 16;
167 pub const IV_TIME_OFFSET: usize = 0;
169 pub const IV_TIME_SIZE: usize = 8;
171 pub const IV_SERVER_ID_OFFSET: usize = 8;
173 pub const IV_SERVER_ID_SIZE: usize = 8;
175 pub const SIGNATURE_SIZE: usize = 4;
177 pub const PAYLOAD_BASE: usize = Crypto::IV_BASE + Crypto::IV_SIZE;
179 pub const OVERHEAD_SIZE: usize = Crypto::IV_SIZE + Crypto::SIGNATURE_SIZE;
181
182 #[inline]
194 pub fn decode<T>(&self, data: T) -> Result<Vec<u8>, CryptoError>
195 where
196 T: AsRef<[u8]>,
197 {
198 URL_SAFE
199 .decode(data)
200 .map(|v| v.to_vec())
201 .map_err(|e| e.into())
202 }
203
204 #[inline]
216 pub fn encode<T>(&self, data: T) -> String
217 where
218 T: AsRef<[u8]>,
219 {
220 URL_SAFE.encode(data)
221 }
222
223 #[inline]
241 pub fn decrypt(&self, cipher_data: &[u8]) -> Result<Vec<u8>, CryptoError> {
242 if cipher_data.len() < Self::OVERHEAD_SIZE {
243 return Err(CryptoError::DataTooShort);
244 }
245
246 let mut data = cipher_data.to_vec();
247 let data_size = data.len();
248
249 self.xor_payload(&mut data)?;
250
251 let confirmation_signature = self.hmac_signature(&data)?;
252 let integrity_signature = self.read_i32(&data, data_size - Self::SIGNATURE_SIZE);
253 self.write_i32(
254 &mut data,
255 data_size - Self::SIGNATURE_SIZE,
256 confirmation_signature,
257 );
258
259 if confirmation_signature != integrity_signature {
260 return Err(CryptoError::InvalidSign);
261 }
262
263 Ok(data)
264 }
265
266 #[inline]
279 pub fn encrypt(&self, plain_data: &[u8]) -> Result<Vec<u8>, CryptoError> {
280 if plain_data.len() < Self::OVERHEAD_SIZE {
281 return Err(CryptoError::DataTooShort);
282 }
283
284 let mut data = plain_data.to_vec();
285 let data_size = data.len();
286 let signature = self.hmac_signature(&data)?;
287 self.write_i32(&mut data, data_size - Self::SIGNATURE_SIZE, signature);
288
289 self.xor_payload(&mut data)?;
290
291 Ok(data)
292 }
293
294 #[inline]
311 pub fn package<T>(&self, payload: T, iv: Option<&[u8]>) -> Result<String, CryptoError>
312 where
313 T: AsRef<[u8]>,
314 {
315 let mut out = Vec::new();
316 self.package_to(payload, iv, &mut out)?;
317 Ok(String::from_utf8(out).expect("base64 output is valid UTF-8"))
319 }
320
321 #[inline]
338 pub fn unpackage<T>(&self, data: T) -> Result<Vec<u8>, CryptoError>
339 where
340 T: AsRef<[u8]>,
341 {
342 let mut out = Vec::new();
343 self.unpackage_to(data, &mut out)?;
344 Ok(out)
345 }
346
347 #[inline]
367 pub fn package_to<T, W>(
368 &self,
369 payload: T,
370 iv: Option<&[u8]>,
371 out: &mut W,
372 ) -> Result<(), CryptoError>
373 where
374 T: AsRef<[u8]>,
375 W: Write,
376 {
377 let payload = payload.as_ref();
378 let mut pkg = self.init_plain_data(payload.len(), iv)?;
379 self.set_payload(&mut pkg, payload)?;
380 let encrypted = self.encrypt(&pkg)?;
381 out.write_all(URL_SAFE.encode(&encrypted).as_bytes())?;
382 Ok(())
383 }
384
385 #[inline]
404 pub fn unpackage_to<T, W>(&self, data: T, out: &mut W) -> Result<(), CryptoError>
405 where
406 T: AsRef<[u8]>,
407 W: Write,
408 {
409 let decoded = self.decode(data)?;
410 let decrypted = self.decrypt(&decoded)?;
411 let payload = self.payload(&decrypted).ok_or(CryptoError::DataTooShort)?;
412 out.write_all(payload)?;
413 Ok(())
414 }
415
416 #[inline]
433 pub fn create_init_vector(&self, timestamp: OffsetDateTime, server_id: i64) -> Vec<u8> {
434 let timestamp = (timestamp.unix_timestamp_nanos() / 1_000) as i64; let mut iv = vec![0; Self::IV_SIZE];
436 self.write_i64(&mut iv, Self::IV_TIME_OFFSET, timestamp);
437 self.write_i64(&mut iv, Self::IV_SERVER_ID_OFFSET, server_id);
438 iv
439 }
440
441 #[inline]
458 pub fn timestamp(&self, data: &[u8]) -> Option<OffsetDateTime> {
459 if data.len() < Self::IV_SIZE {
460 return None;
461 }
462 let ts = self.read_i64(data, Self::IV_BASE + Self::IV_TIME_OFFSET);
463 Some(
464 UNIX_EPOCH
465 .checked_add(Duration::microseconds(ts))
466 .unwrap_or(UNIX_EPOCH),
467 )
468 }
469
470 #[inline]
486 pub fn server_id(&self, data: &[u8]) -> Option<i64> {
487 if data.len() < Self::IV_SIZE {
488 return None;
489 }
490 Some(self.read_i64(data, Self::IV_BASE + Self::IV_SERVER_ID_OFFSET))
491 }
492
493 #[inline]
509 pub fn payload<'a>(&self, data: &'a [u8]) -> Option<&'a [u8]> {
510 if data.len() < Self::OVERHEAD_SIZE {
511 return None;
512 }
513 Some(&data[Self::PAYLOAD_BASE..data.len() - Self::SIGNATURE_SIZE])
514 }
515
516 #[inline]
534 pub fn init_plain_data(
535 &self,
536 payload_size: usize,
537 iv: Option<&[u8]>,
538 ) -> Result<Vec<u8>, CryptoError> {
539 let mut plain_data = vec![0; Self::OVERHEAD_SIZE + payload_size];
540 if let Some(iv) = iv {
541 plain_data[Self::IV_BASE..Self::IV_BASE + Self::IV_SIZE].copy_from_slice(iv);
542 } else {
543 let now = (OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000) as i64;
544 self.write_i64(&mut plain_data, Self::IV_TIME_OFFSET, now);
545 self.write_i64(
546 &mut plain_data,
547 Self::IV_SERVER_ID_OFFSET,
548 rand::random::<i64>(),
549 );
550 }
551
552 Ok(plain_data)
553 }
554
555 #[inline]
572 pub fn set_payload(&self, plain_data: &mut [u8], payload: &[u8]) -> Result<(), CryptoError> {
573 if payload.len() != plain_data.len() - Self::OVERHEAD_SIZE {
574 return Err(CryptoError::PayloadSizeMismatch);
575 }
576 plain_data[Self::PAYLOAD_BASE..Self::PAYLOAD_BASE + payload.len()].copy_from_slice(payload);
577 Ok(())
578 }
579
580 #[inline]
581 fn read_i32(&self, data: &[u8], offset: usize) -> i32 {
582 BigEndian::read_i32(&data[offset..offset + 4])
583 }
584
585 #[inline]
586 fn read_i64(&self, data: &[u8], offset: usize) -> i64 {
587 BigEndian::read_i64(&data[offset..offset + 8])
588 }
589
590 #[inline]
591 fn write_i32(&self, data: &mut [u8], offset: usize, value: i32) {
592 BigEndian::write_i32(&mut data[offset..offset + 4], value);
593 }
594
595 #[inline]
596 fn write_i64(&self, data: &mut [u8], offset: usize, value: i64) {
597 BigEndian::write_i64(&mut data[offset..offset + 8], value);
598 }
599
600 #[inline]
601 fn xor_payload(&self, data: &mut [u8]) -> Result<(), CryptoError> {
602 let iv: &[u8; 16] = &data[Self::IV_BASE..Self::IV_BASE + Self::IV_SIZE]
603 .try_into()
604 .map_err(|_| CryptoError::InvalidInitVector)?;
605
606 let mut cipher = Aes256Ctr64BE::new(&self.keys.encryption_key.into(), iv.into());
607 let data_size = data.len();
608 cipher.apply_keystream(&mut data[Self::PAYLOAD_BASE..data_size - Self::SIGNATURE_SIZE]);
609
610 Ok(())
611 }
612
613 #[inline]
614 fn hmac_signature(&self, data: &[u8]) -> Result<i32, CryptoError> {
615 let mut mac = HmacSha256::new_from_slice(&self.keys.integrity_key)
616 .map_err(|_| CryptoError::InvalidKey)?;
617
618 mac.update(&data[Self::PAYLOAD_BASE..data.len() - Self::SIGNATURE_SIZE]);
619 mac.update(&data[Self::IV_BASE..Self::IV_BASE + Self::IV_SIZE]);
620
621 let b = mac.finalize().into_bytes();
622
623 Ok(self.read_i32(&b, 0))
624 }
625}
626
627#[cfg(test)]
628mod tests {
629 use super::*;
630 use base64::prelude::*;
631
632 static TEST_ENCRYPTION_KEY: &str = "sIxwz7yw62yrfoLGt12lIHKuYrK/S5kLuApI2BQe7Ac=";
633 static TEST_INTEGRITY_KEY: &str = "v3fsVcMBMMHYzRhi7SpM0sdqwzvAxM6KPTu9OtVod5I=";
634
635 fn create_keys() -> Keys {
636 Keys::new(
637 &BASE64_STANDARD.decode(TEST_ENCRYPTION_KEY).unwrap(),
638 &BASE64_STANDARD.decode(TEST_INTEGRITY_KEY).unwrap(),
639 )
640 .unwrap()
641 }
642
643 #[test]
644 fn test_decode() {
645 let crypto = Crypto::new(create_keys());
646 let encoded = "aGVsbG8sIHdvcmxk";
647 let decoded = crypto.decode(encoded).unwrap();
648 assert_eq!(decoded, b"hello, world");
649 }
650
651 #[test]
652 fn test_encode() {
653 let crypto = Crypto::new(create_keys());
654 let data = b"hello, world";
655 let encoded = crypto.encode(data);
656 assert_eq!(encoded, "aGVsbG8sIHdvcmxk");
657 }
658
659 #[test]
660 fn test_decrypt() {
661 let crypto = Crypto::new(create_keys());
662 let timestamp = OffsetDateTime::UNIX_EPOCH + Duration::seconds(1);
663 let iv = crypto.create_init_vector(timestamp, 123456789);
664 let payload = "https://example.com".as_bytes();
665
666 let mut plain_data = crypto.init_plain_data(payload.len(), Some(&iv)).unwrap();
667 crypto.set_payload(&mut plain_data, payload).unwrap();
668 let encrypted_data = crypto.encrypt(&plain_data).unwrap();
669
670 assert_eq!(crypto.timestamp(&iv), Some(timestamp));
671 assert_eq!(crypto.server_id(&iv), Some(123456789));
672 assert_eq!(
673 crypto.payload(&encrypted_data).unwrap().len(),
674 payload.len()
675 );
676 assert_ne!(crypto.payload(&encrypted_data), Some(payload));
677
678 let decrypted_data = crypto.decrypt(&encrypted_data).unwrap();
679 assert_eq!(crypto.timestamp(&decrypted_data), Some(timestamp));
680 assert_eq!(crypto.server_id(&decrypted_data), Some(123456789));
681 assert_eq!(crypto.payload(&decrypted_data), Some(payload));
682
683 let mut encrypted_data_invalid_sign = encrypted_data.clone();
684 crypto.write_i32(
685 &mut encrypted_data_invalid_sign,
686 encrypted_data.len() - Crypto::SIGNATURE_SIZE,
687 123456789,
688 );
689 assert!(matches!(
690 crypto.decrypt(&encrypted_data_invalid_sign),
691 Err(CryptoError::InvalidSign)
692 ));
693 assert_ne!(crypto.payload(&encrypted_data_invalid_sign), Some(payload))
694 }
695
696 #[test]
697 fn test_create_init_vector() {
698 let crypto = Crypto::new(create_keys());
699 let timestamp = OffsetDateTime::UNIX_EPOCH + Duration::seconds(1);
700 let iv = crypto.create_init_vector(timestamp, 123456789);
701 assert_eq!(iv.len(), Crypto::IV_SIZE);
702 assert_eq!(crypto.read_i64(&iv, Crypto::IV_TIME_OFFSET), 1_000_000);
703 assert_eq!(crypto.read_i64(&iv, Crypto::IV_SERVER_ID_OFFSET), 123456789);
704 assert_eq!(crypto.timestamp(&iv), Some(timestamp));
705 assert_eq!(crypto.server_id(&iv), Some(123456789));
706 }
707
708 #[test]
709 fn test_init_plain_data() {
710 let crypto = Crypto::new(create_keys());
711 let payload = "https://example.com".as_bytes();
712
713 let mut plain_data = crypto.init_plain_data(payload.len(), None).unwrap();
714 crypto.set_payload(&mut plain_data, payload).unwrap();
715
716 assert_eq!(plain_data.len(), Crypto::OVERHEAD_SIZE + payload.len());
717 assert_eq!(crypto.payload(&plain_data), Some(payload));
718 }
719
720 #[test]
721 fn test_init_plain_data_empty_payload() {
722 let crypto = Crypto::new(create_keys());
723 let payload = "".as_bytes();
724
725 let mut plain_data = crypto.init_plain_data(0, None).unwrap();
726 crypto.set_payload(&mut plain_data, payload).unwrap();
727 assert_eq!(crypto.payload(&plain_data), Some(payload));
728 }
729
730 #[test]
731 fn test_package_unpackage() {
732 let crypto = Crypto::new(create_keys());
733 let payload = b"Hello, world!".as_slice();
734
735 let encoded = crypto.package(payload, None).unwrap();
736 assert_ne!(encoded, "");
737
738 let decoded = crypto.unpackage(&encoded).unwrap();
739 assert_eq!(decoded, payload);
740 }
741
742 #[test]
743 fn test_package_unpackage_with_iv() {
744 let crypto = Crypto::new(create_keys());
745 let timestamp = OffsetDateTime::UNIX_EPOCH + Duration::seconds(1);
746 let iv = crypto.create_init_vector(timestamp, 123456789);
747 let payload = b"https://example.com".as_slice();
748
749 let encoded = crypto.package(payload, Some(&iv)).unwrap();
750
751 let decoded = crypto.decode(encoded.as_bytes()).unwrap();
753 assert_eq!(crypto.timestamp(&decoded), Some(timestamp));
754 assert_eq!(crypto.server_id(&decoded), Some(123456789));
755
756 let recovered = crypto.unpackage(&encoded).unwrap();
757 assert_eq!(recovered, payload);
758 }
759
760 #[test]
761 fn test_package_unpackage_empty_payload() {
762 let crypto = Crypto::new(create_keys());
763 let payload = b"".as_slice();
764
765 let encoded = crypto.package(payload, None).unwrap();
766 let recovered = crypto.unpackage(&encoded).unwrap();
767 assert_eq!(recovered, payload);
768 }
769
770 #[test]
771 fn test_unpackage_tampered_signature() {
772 let crypto = Crypto::new(create_keys());
773 let encoded = crypto.package(b"Hello, world!", None).unwrap();
774
775 let mut bytes = crypto.decode(encoded.as_bytes()).unwrap();
776 let last = bytes.len() - Crypto::SIGNATURE_SIZE;
777 crypto.write_i32(&mut bytes, last, 123456789);
778
779 let tampered = crypto.encode(&bytes);
780 assert!(matches!(
781 crypto.unpackage(&tampered),
782 Err(CryptoError::InvalidSign)
783 ));
784 }
785
786 #[test]
787 fn test_package_to_matches_package() {
788 let crypto = Crypto::new(create_keys());
789 let timestamp = OffsetDateTime::UNIX_EPOCH + Duration::seconds(1);
790 let iv = crypto.create_init_vector(timestamp, 123456789);
791 let payload = b"https://example.com".as_slice();
792
793 let encoded_alloc = crypto.package(payload, Some(&iv)).unwrap();
795
796 let mut buf = Vec::new();
797 crypto.package_to(payload, Some(&iv), &mut buf).unwrap();
798 assert_eq!(buf, encoded_alloc.as_bytes());
799 }
800
801 #[test]
802 fn test_package_to_unpackage_to_roundtrip() {
803 let crypto = Crypto::new(create_keys());
804 let payload = b"Hello, world!".as_slice();
805
806 let mut enc_buf = Vec::new();
807 crypto.package_to(payload, None, &mut enc_buf).unwrap();
808
809 let mut dec_buf = Vec::new();
810 crypto.unpackage_to(&enc_buf, &mut dec_buf).unwrap();
811 assert_eq!(dec_buf, payload);
812 }
813
814 #[test]
815 fn test_package_to_appends_and_preserves_existing() {
816 let crypto = Crypto::new(create_keys());
817 let payload = b"Hello".as_slice();
818
819 let mut buf = b"prefix".to_vec();
820 let prefix_len = buf.len();
821 crypto.package_to(payload, None, &mut buf).unwrap();
822
823 assert_eq!(&buf[..prefix_len], b"prefix");
825 assert!(buf.len() > prefix_len);
826
827 let mut dec_buf = Vec::new();
828 crypto
829 .unpackage_to(&buf[prefix_len..], &mut dec_buf)
830 .unwrap();
831 assert_eq!(dec_buf, payload);
832 }
833
834 #[test]
835 fn test_package_to_empty_payload() {
836 let crypto = Crypto::new(create_keys());
837
838 let mut enc_buf = Vec::new();
839 crypto.package_to(b"", None, &mut enc_buf).unwrap();
840
841 let mut dec_buf = Vec::new();
842 crypto.unpackage_to(&enc_buf, &mut dec_buf).unwrap();
843 assert_eq!(dec_buf, b"");
844 }
845
846 #[test]
847 fn test_unpackage_to_tampered_signature() {
848 let crypto = Crypto::new(create_keys());
849
850 let mut enc_buf = Vec::new();
851 crypto.package_to(b"Hello", None, &mut enc_buf).unwrap();
852
853 let mut raw = crypto.decode(&enc_buf).unwrap();
856 let last = raw.len() - Crypto::SIGNATURE_SIZE;
857 crypto.write_i32(&mut raw, last, 123456789);
858 let tampered = crypto.encode(&raw);
859
860 let mut dec_buf = Vec::new();
861 assert!(matches!(
862 crypto.unpackage_to(tampered.as_bytes(), &mut dec_buf),
863 Err(CryptoError::InvalidSign)
864 ));
865 }
866
867 #[test]
868 fn test_package_to_with_non_vec_writer() {
869 let crypto = Crypto::new(create_keys());
872 let payload = b"Hello, world!".as_slice();
873
874 let mut writer = std::io::BufWriter::new(Vec::<u8>::new());
875 crypto.package_to(payload, None, &mut writer).unwrap();
876 let encoded = writer.into_inner().unwrap();
877
878 let mut dec_buf = Vec::new();
879 crypto.unpackage_to(&encoded, &mut dec_buf).unwrap();
880 assert_eq!(dec_buf, payload);
881 }
882
883 #[test]
884 fn test_unpackage_to_appends_to_existing_buffer() {
885 let crypto = Crypto::new(create_keys());
888 let payload = b"Hello".as_slice();
889 let encoded = crypto.package(payload, None).unwrap();
890
891 let mut buf = b"prefix".to_vec();
892 let prefix_len = buf.len();
893 crypto.unpackage_to(&encoded, &mut buf).unwrap();
894
895 assert_eq!(&buf[..prefix_len], b"prefix");
896 assert_eq!(&buf[prefix_len..], payload);
897 }
898}