More flexible PCS support

TaeHagen · TaeHagen · commit 077fa566d668 · 2025-08-28T20:29:39.000-04:00
diff --git a/cloudkit-derive/src/lib.rs b/cloudkit-derive/src/lib.rs
@@ -12,6 +12,7 @@ use syn::{parse_macro_input, Data, DeriveInput, LitStr};
 struct CloudKitRecordAttributes {
     r#type: String,
     encrypted: Flag,
+    rename_all: Option<String>,
 }
 
 #[derive(deluxe::ExtractAttributes)]
@@ -22,11 +23,29 @@ struct CloudKitAttributes {
     unencrypted: Flag,
 }
 
+fn snake_to_camel(s: &str) -> String {
+    let mut camel = String::new();
+    let mut upper_next = false;
+
+    for c in s.chars() {
+        if c == '_' {
+            upper_next = true;
+        } else if upper_next {
+            camel.push_str(&c.to_uppercase().to_string());
+            upper_next = false;
+        } else {
+            camel.push(c);
+        }
+    }
+
+    camel
+}
+
 #[proc_macro_derive(CloudKitRecord, attributes(cloudkit_record, cloudkit))]
 pub fn cloudkitrecord_derive(input: TokenStream) -> TokenStream {
     let mut input = parse_macro_input!(input as DeriveInput);
 
-    let CloudKitRecordAttributes { r#type, encrypted: record_encrypted } = deluxe::extract_attributes(&mut input).unwrap();
+    let CloudKitRecordAttributes { r#type, encrypted: record_encrypted, rename_all } = deluxe::extract_attributes(&mut input).unwrap();
 
     let name = input.ident;
 
@@ -48,7 +67,18 @@ pub fn cloudkitrecord_derive(input: TokenStream) -> TokenStream {
         }
         
         let ident = field.ident.unwrap();
-        let name = rename.unwrap_or_else(|| ident.to_string());
+        let name = rename.unwrap_or_else(|| {
+            if let Some(rename_all) = &rename_all {
+                let name = ident.to_string();
+                return if rename_all == "camelCase" {
+                    snake_to_camel(&name)
+                } else {
+                    panic!("unknown rename {}", rename_all)
+                }
+            }
+
+            ident.to_string()
+        });
         let name_lit = LitStr::new(&name, Span::call_site());
         if is_encrypted {
             fields.push(quote! {
diff --git a/cloudkit-proto/src/cloudkit.proto b/cloudkit-proto/src/cloudkit.proto
@@ -906,6 +906,7 @@ message RequestOperation {
 
         message EncryptedValue {
           optional int64 signedValue = 3;
+          optional Date dateValue = 5;
           optional string stringValue = 6;
         }
   
@@ -983,6 +984,17 @@ message RequestOperation {
       optional uint32 permission = 15;
       optional bytes pcsKey = 24;
   }
+
+  // the abuse of the cloudkit file continues, this should really go in a PCS file, but it only has two messages...
+  message ProtoPCSPrivateKey {
+    required bytes key = 1;
+    optional bytes public = 2;
+  }
+
+  message ProtoPCSKey {
+    required ProtoPCSPrivateKey encryptionKey = 1;
+    optional ProtoPCSPrivateKey signingKey = 2;
+  }
   
   message RequestedFields {
      repeated Record.Field.Identifier fields = 1;
diff --git a/cloudkit-proto/src/lib.rs b/cloudkit-proto/src/lib.rs
@@ -205,6 +205,30 @@ impl CloudKitEncryptedValue for String {
     }
 }
 
+impl CloudKitEncryptedValue for SystemTime {
+    fn to_value_encrypted(&self, encryptor: &impl CloudKitEncryptor, context: &[u8]) -> Option<record::field::Value> {
+        let apple_epoch = SystemTime::UNIX_EPOCH + Duration::from_secs(978307200);
+        let duration = self.duration_since(apple_epoch).expect("Before Apple Epoch?").as_secs_f64();
+
+        Some(record::field::Value {
+            r#type: Some(FieldType::DateType as i32),
+            bytes_value: Some(encryptor.encrypt_data(&record::field::EncryptedValue {
+                date_value: Some(Date { time: Some(duration) }),
+                ..Default::default()
+            }.encode_to_vec(), context)),
+            is_encrypted: Some(true),
+            ..Default::default()
+        })
+    }
+
+    fn from_value_encrypted(value: &record::field::Value, encryptor: &impl CloudKitEncryptor, context: &[u8]) -> Option<Self> {
+        let d = record::field::EncryptedValue::decode(&encryptor.decrypt_data(value.bytes_value.as_ref().unwrap(), context)[..]).unwrap().date_value.unwrap();
+        let secs = d.time.expect("Date misses time??");
+        let apple_epoch = SystemTime::UNIX_EPOCH + Duration::from_secs(978307200);
+        Some(apple_epoch + Duration::from_secs_f64(secs))
+    }
+}
+
 pub fn encode_hex(bytes: &[u8]) -> String {
     use std::fmt::Write;
     let mut s = String::with_capacity(bytes.len() * 2);
diff --git a/src/cloudkit.rs b/src/cloudkit.rs
@@ -42,15 +42,16 @@ pub struct FetchedRecords {
 }
 
 impl FetchedRecords {
-    pub fn get_record<R: CloudKitRecord>(&self, record_id: &str, key: Option<&PCSKey>) -> R {
+    pub fn get_record<R: CloudKitRecord>(&self, record_id: &str, key: Option<&PCSKeys>) -> R {
         self.responses.iter().find_map(|response| {
             let r = response.record_retrieve_response.as_ref().expect("No retrieve response?").record.as_ref().expect("No record?");
             if r.record_identifier.as_ref().expect("No record id?").value.as_ref().expect("No identifier").name.as_ref().expect("No name?") == record_id {                
                 let got_type = r.r#type.as_ref().expect("no TYpe").name.as_ref().expect("No ta");
                 if got_type.as_str() != R::record_type() {
                     panic!("Wrong record type, got {} expected {}", got_type, R::record_type());
                 }
-                Some(R::from_record_encrypted(&r.record_field, key.map(|k| (k, r.record_identifier.as_ref().unwrap()))))
+                let key = key.map(|k| pcs_key_for_record(r, k).expect("PCS key failed"));
+                Some(R::from_record_encrypted(&r.record_field, key.as_ref().map(|k| (k, r.record_identifier.as_ref().unwrap()))))
             } else { None }
         }).expect("No record found?")
     }
@@ -98,6 +99,13 @@ pub trait CloudKitOp {
     }
 }
 
+pub fn pcs_key_for_record(record: &Record, keys: &PCSKeys) -> Result<PCSKey, PushError> {
+    let Some(protection) = &record.protection_info else {
+        return Ok(keys.default_record_key.clone().unwrap())
+    };
+    keys.decode_record_protection(protection)
+}
+
 pub struct UploadAssetOperation(pub cloudkit_proto::AssetUploadTokenRetrieveRequest);
 impl CloudKitOp for UploadAssetOperation {
     type Response = cloudkit_proto::AssetUploadTokenRetrieveResponse;
@@ -178,15 +186,35 @@ impl CloudKitOp for SaveRecordOperation {
 }
 
 impl SaveRecordOperation {
-    pub fn new<R: CloudKitRecord>(id: RecordIdentifier, record: R, key: Option<&PCSKey>, update: bool) -> Self {
+    pub fn new<R: CloudKitRecord>(id: RecordIdentifier, record: R, key: Option<&PCSKeys>, update: bool) -> Self {
+        let mut pcs_key_id: Option<Vec<u8>> = None;
+        let mut protection_info: Option<ProtectionInfo> = None;
+        let mut pcs_key: Option<PCSKey> = None;
+
+        if let Some(key) = key {
+            if let Some(global) = &key.default_record_key {
+                pcs_key_id = Some(global.key_id().unwrap()[..4].to_vec());
+                pcs_key = Some(global.clone());
+            } else {
+                // create a key for this record
+                let record_protection = PCSShareProtection::create(&key.zone_keys[0], &[]).unwrap();
+                let der = rasn::der::encode(&record_protection).unwrap();
+                protection_info = Some(ProtectionInfo { 
+                    protection_info_tag: Some(encode_hex(&sha1(&der)).to_uppercase()),
+                    protection_info: Some(der), 
+                });
+            }
+        }
+
         Self(cloudkit_proto::RecordSaveRequest {
             record: Some(cloudkit_proto::Record {
                 record_identifier: Some(id.clone()),
                 r#type: Some(cloudkit_proto::record::Type {
                     name: Some(R::record_type().to_string())
                 }),
-                record_field: record.to_record_encrypted(key.map(|k| (k, &id))),
-                pcs_key: key.map(|k| k.key_id().unwrap()[..4].to_vec()),
+                record_field: record.to_record_encrypted(pcs_key.as_ref().map(|k| (k, &id))),
+                pcs_key: pcs_key_id,
+                protection_info,
                 ..Default::default()
             }),
             merge: Some(true),
@@ -203,14 +231,15 @@ pub struct FetchedRecord {
 }
 
 impl FetchedRecord {
-    pub fn get_record<R: CloudKitRecord>(&self, key: Option<&PCSKey>) -> R {
+    pub fn get_record<R: CloudKitRecord>(&self, key: Option<&PCSKeys>) -> R {
         let r = self.response.record_retrieve_response.as_ref().expect("No retrieve response?").record.as_ref().expect("No record?");
         
         let got_type = r.r#type.as_ref().expect("no TYpe").name.as_ref().expect("No ta");
         if got_type.as_str() != R::record_type() {
             panic!("Wrong record type, got {} expected {}", got_type, R::record_type());
         }
-        R::from_record_encrypted(&r.record_field, key.map(|k| (k, r.record_identifier.as_ref().unwrap())))
+        let key = key.map(|k| pcs_key_for_record(r, k).expect("no PCS key"));
+        R::from_record_encrypted(&r.record_field, key.as_ref().map(|k| (k, r.record_identifier.as_ref().unwrap())))
     }
 
     pub fn get_id(&self) -> String {
@@ -546,15 +575,17 @@ impl CloudKitOp for ZoneSaveOperation {
 }
 
 impl ZoneSaveOperation {
-    pub fn new(zone: RecordZoneIdentifier, pcs_key: Option<&CompactECKey<Private>>) -> Result<Self, PushError> {
+    pub fn new(zone: RecordZoneIdentifier, pcs_key: Option<&CompactECKey<Private>>, with_record: bool) -> Result<Self, PushError> {
         let mut protection_info: Option<ProtectionInfo> = None;
         let mut record_protection_info: Option<ProtectionInfo> = None;
         if let Some(pcs_key) = pcs_key {
             let zone_key = CompactECKey::new()?;
-            let record_protection = PCSShareProtection::create(&zone_key, &[])?;
-            let main_protection = PCSShareProtection::create(pcs_key, &[zone_key])?;
-
-            record_protection_info = Some(ProtectionInfo { protection_info: Some(rasn::der::encode(&record_protection).unwrap()), protection_info_tag: None });
+            let main_protection = PCSShareProtection::create(pcs_key, &[zone_key.clone()])?;
+            
+            if with_record {
+                let record_protection = PCSShareProtection::create(&zone_key, &[])?;
+                record_protection_info = Some(ProtectionInfo { protection_info: Some(rasn::der::encode(&record_protection).unwrap()), protection_info_tag: None });
+            }
             let main_encoded = rasn::der::encode(&main_protection).unwrap();
             protection_info = Some(ProtectionInfo {
                 protection_info_tag: Some(encode_hex(&sha1(&main_encoded)).to_uppercase()),
@@ -715,11 +746,38 @@ pub struct QueryResult<T: CloudKitRecord> {
     pub result: T,
 }
 
+#[derive(Clone)]
+pub struct PCSKeys {
+    zone_keys: Vec<CompactECKey<Private>>,
+    default_record_key: Option<PCSKey>,
+}
+
+impl PCSKeys {
+    fn new(keys: Vec<CompactECKey<Private>>) -> Self {
+        Self {
+            zone_keys: keys,
+            default_record_key: None,
+        }
+    }
+
+    fn decode_record_protection(&self, protection: &ProtectionInfo) -> Result<PCSKey, PushError> {
+        let record_protection: PCSShareProtection = rasn::der::decode(protection.protection_info()).expect("Bad record protection?");
+        let mut big_num = BigNumContext::new()?;
+        let record_key = CompactECKey::decompress(record_protection.decode_key_public()?.try_into().expect("Decode key not compact!"));
+
+        let item = self.zone_keys.iter().find(|k| matches!(record_key.public_key().eq(&record_key.group(), &k.public_key(), &mut big_num), Ok(true))).expect("Record key not found!");
+
+        let (key, _record_keys) = record_protection.decode(item).unwrap();
+
+        Ok(key)
+    }
+}
+
 pub struct CloudKitOpenContainer<'t, T: AnisetteProvider> {
     container: CloudKitContainer<'t>,
     pub user_id: String,
     pub client: Arc<CloudKitClient<T>>,
-    pub keys: Mutex<HashMap<String, PCSKey>>,
+    pub keys: Mutex<HashMap<String, PCSKeys>>,
 }
 
 impl<'t, T: AnisetteProvider> Deref for CloudKitOpenContainer<'t, T> {
@@ -744,14 +802,14 @@ impl<'t, T: AnisetteProvider> CloudKitOpenContainer<'t, T> {
         }
     }
 
-    pub async fn get_zone_encryption_config(&self, zone: &cloudkit_proto::RecordZoneIdentifier, client: &KeychainClient<T>, service: &PCSService<'_>) -> Result<PCSKey, PushError> {
+    pub async fn get_zone_encryption_config(&self, zone: &cloudkit_proto::RecordZoneIdentifier, client: &KeychainClient<T>, pcs_service: &PCSService<'_>) -> Result<PCSKeys, PushError> {
         let mut cached_keys = self.keys.lock().await;
         let zone_name = zone.value.as_ref().unwrap().name().to_string();
         if let Some(key) = cached_keys.get(&zone_name) {
             return Ok(key.clone());
         }
         
-        client.sync_keychain(&[&service.zone, "ProtectedCloudStorage"]).await?;
+        client.sync_keychain(&[&pcs_service.zone, "ProtectedCloudStorage"]).await?;
 
         let zone = match self.perform(&CloudKitSession::new(), FetchZoneOperation::new(zone.clone())).await {
             Ok(data) => data.target_zone.unwrap(),
@@ -764,9 +822,9 @@ impl<'t, T: AnisetteProvider> CloudKitOpenContainer<'t, T> {
                 }),
                 ..
             })) => {
-                let service = PCSPrivateKey::get_service_key(client, service, self.client.config.as_ref()).await?;
+                let service = PCSPrivateKey::get_service_key(client, pcs_service, self.client.config.as_ref()).await?;
 
-                let request = ZoneSaveOperation::new(zone.clone(), Some(&service.key())).unwrap();
+                let request = ZoneSaveOperation::new(zone.clone(), Some(&service.key()), pcs_service.global_record).unwrap();
                 let zone = request.0.clone().zone.unwrap();
                 self.perform(&CloudKitSession::new(), request).await.unwrap();
                 zone
@@ -777,20 +835,17 @@ impl<'t, T: AnisetteProvider> CloudKitOpenContainer<'t, T> {
 
         let data = client.state.read().await;
 
-        let (_parent_key, keys) = zone_protection.decrypt_with_keychain(&data)?;
+        let (_parent_key, keys) = zone_protection.decrypt_with_keychain(&data, pcs_service)?;
 
-        let record_protection: PCSShareProtection = rasn::der::decode(zone.record_protection_info.as_ref().unwrap().protection_info()).expect("Bad record protection?");
-
-        let mut big_num = BigNumContext::new()?;
-        let record_key = CompactECKey::decompress(record_protection.decode_key_public()?.try_into().expect("Decode key not compact!"));
-
-        let item = keys.iter().find(|k| matches!(record_key.public_key().eq(&record_key.group(), &k.public_key(), &mut big_num), Ok(true))).expect("Record key not found!");
-
-        let (key, _record_keys) = record_protection.decode(item).unwrap();
-
-        cached_keys.insert(zone_name, key.clone());
+        let mut keys = PCSKeys::new(keys);
+        
+        if let Some(record_protection_info) = &zone.record_protection_info {
+            keys.default_record_key = Some(keys.decode_record_protection(record_protection_info)?);
+        }
+        
+        cached_keys.insert(zone_name, keys.clone());
 
-        Ok(key)
+        Ok(keys)
     }
 
     pub fn build_request<Op: CloudKitOp>(&self, operation: &Op, config: &dyn OSConfig, is_first: bool, uuid: String, isolation_level: IsolationLevel) -> Vec<u8> {
diff --git a/src/imessage/cloud_messages.rs b/src/imessage/cloud_messages.rs
@@ -26,7 +26,7 @@ use cloudkit_proto::RecordIdentifier;
 use log::info;
 use uuid::Uuid;
 use crate::cloud_messages::cloudmessagesp::{ChatProto, MessageProto, MessageProto2, MessageProto3, MessageProto4};
-use crate::cloudkit::{record_identifier, CloudKitSession, CloudKitUploadRequest, DeleteRecordOperation, FetchRecordChangesOperation, FetchRecordOperation, FetchedRecords, SaveRecordOperation, ZoneDeleteOperation, ALL_ASSETS, NO_ASSETS};
+use crate::cloudkit::{pcs_key_for_record, record_identifier, CloudKitSession, CloudKitUploadRequest, DeleteRecordOperation, FetchRecordChangesOperation, FetchRecordOperation, FetchedRecords, SaveRecordOperation, ZoneDeleteOperation, ALL_ASSETS, NO_ASSETS};
 use crate::mmcs::{prepare_put_v2, PreparedPut};
 use crate::pcs::{get_boundary_key, PCSKey, PCSService};
 use bitflags::bitflags;
@@ -43,6 +43,8 @@ pub const MESSAGES_SERVICE: PCSService = PCSService {
     zone: "Engram",
     r#type: 55,
     keychain_type: 55,
+    v2: false,
+    global_record: true,
 };
 
 pub mod cloudmessagesp {
@@ -477,7 +479,7 @@ impl<P: AnisetteProvider> CloudMessagesClient<P> {
             };
             if record.r#type.as_ref().unwrap().name() != T::record_type() { continue }
 
-            let item = T::from_record_encrypted(&record.record_field, Some((&key, record.record_identifier.as_ref().unwrap())));
+            let item = T::from_record_encrypted(&record.record_field, Some((&pcs_key_for_record(&record, &key)?, record.record_identifier.as_ref().unwrap())));
 
             results.insert(identifier, Some(item));
         }
diff --git a/src/keychain.rs b/src/keychain.rs
@@ -913,7 +913,7 @@ impl KeychainClientState {
     }
 }
 
-const KEYCHAIN_ZONES: &[&str] = &[
+pub const KEYCHAIN_ZONES: &[&str] = &[
     "AutoUnlock",
     "SecureObjectSync",
     "SE-PTC",
diff --git a/src/pcs.rs b/src/pcs.rs

Original file line number	Diff line number	Diff line change
`@@ -913,7 +913,7 @@ impl KeychainClientState {`
`913`	`913`	`}`
`914`	`914`	`}`
`915`	`915`
`916`		`-const KEYCHAIN_ZONES: &[&str] = &[`
	`916`	`+pub const KEYCHAIN_ZONES: &[&str] = &[`
`917`	`917`	`"AutoUnlock",`
`918`	`918`	`"SecureObjectSync",`
`919`	`919`	`"SE-PTC",`