PCS fixes

TaeHagen · TaeHagen · commit 5d5954442ba2 · 2025-08-31T16:41:15.000-04:00
diff --git a/src/cloudkit.rs b/src/cloudkit.rs
@@ -101,9 +101,12 @@ 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())
+        let Some(pcskey) = &record.pcs_key else { panic!("No PCS Key??") };
+        let keys = keys.default_record_keys.iter().find(|i| i.key_id().ok().map(|id| pcskey == &id[..pcskey.len()]).unwrap_or(false));
+        
+        return Ok(keys.ok_or(PushError::PCSRecordKeyMissing)?.clone())
     };
-    keys.decode_record_protection(protection)
+    Ok(keys.decode_record_protection(protection)?.remove(0))
 }
 
 pub struct UploadAssetOperation(pub cloudkit_proto::AssetUploadTokenRetrieveRequest);
@@ -196,7 +199,7 @@ impl SaveRecordOperation {
             protection_info_tag: Some(tag.clone()),
             protection_info: Some(der), 
         });
-        let pcs_key = key.decode_record_protection(protection_info.as_ref().unwrap()).expect("Failed to decode record protection");
+        let pcs_key = key.decode_record_protection(protection_info.as_ref().unwrap()).expect("Failed to decode record protection").remove(0);
 
         (Self(cloudkit_proto::RecordSaveRequest {
             record: Some(cloudkit_proto::Record {
@@ -222,8 +225,8 @@ impl SaveRecordOperation {
                 r#type: Some(cloudkit_proto::record::Type {
                     name: Some(R::record_type().to_string())
                 }),
-                record_field: record.to_record_encrypted(key.map(|k| (k.default_record_key.as_ref().expect("No default record key?"), &id))),
-                pcs_key: key.map(|k| k.default_record_key.as_ref().expect("No default record key?").key_id().unwrap()[..4].to_vec()),
+                record_field: record.to_record_encrypted(key.map(|k| (k.default_record_keys.first().expect("No default record key?"), &id))),
+                pcs_key: key.map(|k| k.default_record_keys.first().expect("No default record key?").key_id().unwrap()[..4].to_vec()),
                 ..Default::default()
             }),
             merge: Some(true),
@@ -768,13 +771,13 @@ pub struct QueryResult<T: CloudKitRecord> {
 pub struct PCSKeys {
     zone_keys: Vec<CompactECKey<Private>>,
     zone_protection_tag: Option<String>,
-    default_record_key: Option<PCSKey>,
+    default_record_keys: Vec<PCSKey>,
     pub record_prot_tag: Option<String>,
 }
 
 impl PCSKeys {
 
-    fn decode_record_protection(&self, protection: &ProtectionInfo) -> Result<PCSKey, PushError> {
+    fn decode_record_protection(&self, protection: &ProtectionInfo) -> Result<Vec<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!"));
@@ -816,6 +819,12 @@ impl<'t, T: AnisetteProvider> CloudKitOpenContainer<'t, T> {
         }
     }
 
+    pub async fn clear_cache_zone_encryption_config(&self, zone: &cloudkit_proto::RecordZoneIdentifier) {
+        let mut cached_keys = self.keys.lock().await;
+        let zone_name = zone.value.as_ref().unwrap().name().to_string();
+        cached_keys.remove(&zone_name);
+    }
+
     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();
@@ -854,14 +863,14 @@ impl<'t, T: AnisetteProvider> CloudKitOpenContainer<'t, T> {
         let mut keys = PCSKeys {
             zone_keys: keys,
             zone_protection_tag: zone.protection_info.as_ref().unwrap().protection_info_tag.clone(),
-            default_record_key: None,
+            default_record_keys: vec![],
             record_prot_tag: if let Some(record_protection_info) = &zone.record_protection_info {
                 record_protection_info.protection_info_tag.clone()
             } else { None },
         };
 
         if let Some(record_protection_info) = &zone.record_protection_info {
-            keys.default_record_key = Some(keys.decode_record_protection(record_protection_info)?);
+            keys.default_record_keys = keys.decode_record_protection(record_protection_info)?;
         }
         
         cached_keys.insert(zone_name, keys.clone());
diff --git a/src/error.rs b/src/error.rs
@@ -177,4 +177,6 @@ pub enum PushError {
     BatchError(Arc<PushError>),
     #[error("Invalid 2fa code!")]
     Bad2FaCode,
+    #[error("PCS record key id not found!")]
+    PCSRecordKeyMissing,
 }
diff --git a/src/imessage/cloud_messages.rs b/src/imessage/cloud_messages.rs
@@ -345,29 +345,52 @@ impl Default for NumOrString {
 #[derive(Serialize, Deserialize, Debug, Clone, Default)]
 #[serde(rename_all = "kebab-case")]
 pub struct MMCSAttachmentMeta {
-    pub mmcs_signature_hex: String,
+    // MMCS attachments
+    pub mmcs_signature_hex: Option<String>,
+    pub mmcs_owner: Option<String>,
+    pub mmcs_url: Option<String>,
+    pub decryption_key: Option<String>,
+
+    // inline attachments
+    pub inline_attachment: Option<String>,
+    pub message_part: Option<String>,
+
     pub file_size: NumOrString,
-    pub decryption_key: String,
     pub uti_type: Option<String>,
-    pub mmcs_owner: String,
     pub mime_type: Option<String>,
-    pub mmcs_url: String,
     pub name: Option<String>,
 }
 
 
 impl Into<Option<MMCSAttachmentMeta>> for &Attachment {
     fn into(self) -> Option<MMCSAttachmentMeta> {
         match &self.a_type {
-            AttachmentType::Inline(_inline) => None,
+            AttachmentType::Inline(_inline) => Some(MMCSAttachmentMeta { 
+                mmcs_signature_hex: None, 
+                decryption_key: None, 
+                mmcs_owner: None, 
+                mmcs_url: None, 
+
+                inline_attachment: Some("ia-0".to_string()),
+                message_part: Some("0".to_string()),
+
+                file_size: NumOrString::Num(_inline.len() as u32), 
+                uti_type: Some(self.uti_type.clone()), 
+                mime_type: Some(self.mime.clone()),
+                name: Some(self.name.clone())
+            }),
             AttachmentType::MMCS(mmcs) => Some(MMCSAttachmentMeta { 
-                mmcs_signature_hex: encode_hex(&mmcs.signature), 
+                mmcs_signature_hex: Some(encode_hex(&mmcs.signature)), 
+                decryption_key: Some(encode_hex(&mmcs.key)), 
+                mmcs_owner: Some(mmcs.object.clone()), 
+                mmcs_url: Some(mmcs.url.clone()), 
+
+                inline_attachment: None,
+                message_part: None,
+
                 file_size: NumOrString::Num(mmcs.size as u32), 
-                decryption_key: encode_hex(&mmcs.key), 
                 uti_type: Some(self.uti_type.clone()), 
-                mmcs_owner: mmcs.object.clone(), 
                 mime_type: Some(self.mime.clone()),
-                mmcs_url: mmcs.url.clone(), 
                 name: Some(self.name.clone())
             })
         }
@@ -399,7 +422,7 @@ pub struct AttachmentMeta {
     #[serde(rename = "ui")]
     pub user_info: Option<MMCSAttachmentMeta>,
     #[serde(rename = "fn")]
-    pub filename: String, //path
+    pub filename: Option<String>, //path
     #[serde(rename = "aui")]
     pub extras: Option<AttachmentMetaExtra>,
     #[serde(rename = "ig")]
@@ -479,7 +502,15 @@ 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((&pcs_key_for_record(&record, &key)?, record.record_identifier.as_ref().unwrap())));
+            let pcskey = match pcs_key_for_record(&record, &key) {
+                Ok(key) => key,
+                Err(PushError::PCSRecordKeyMissing) => {
+                    container.clear_cache_zone_encryption_config(&zone).await;
+                    return Err(PushError::PCSRecordKeyMissing)
+                },
+                Err(e) => return Err(e)
+            };
+            let item = T::from_record_encrypted(&record.record_field, Some((&pcskey, record.record_identifier.as_ref().unwrap())));
 
             results.insert(identifier, Some(item));
         }
@@ -543,6 +574,7 @@ impl<P: AnisetteProvider> CloudMessagesClient<P> {
             ZoneDeleteOperation::new(container.private_zone("chatManateeZone".to_string())),
             ZoneDeleteOperation::new(container.private_zone("messageManateeZone".to_string())),
             ZoneDeleteOperation::new(container.private_zone("attachmentManateeZone".to_string())),
+            ZoneDeleteOperation::new(container.private_zone("chat1ManateeZone".to_string())),
         ], IsolationLevel::Operation).await?;
         Ok(())
     }
diff --git a/src/pcs.rs b/src/pcs.rs
@@ -311,7 +311,7 @@ impl PCSPrivateKey {
     pub fn from_dict(dict: &Dictionary, keychain: &KeychainClientState) -> Self {
         let key = dict.get("v_Data").expect("No dat?").as_data().expect("Not data");
 
-        let decoded: PCSPrivateKey = rasn::der::decode(&key).unwrap();
+        let decoded: PCSPrivateKey = rasn::der::decode(&key).expect("Failed to decode private key!");
 
         if !decoded.verify_with_keychain(keychain, dict.get("atyp").expect("No dat?").as_data().expect("Not data")).unwrap() {
             panic!("PCS Master key verification failed!");
@@ -559,6 +559,8 @@ impl PCSShareProtectionKeySet {
 
 #[derive(AsnType, Encode, Decode)]
 pub struct PCSShareProtectionIdentities {
+    #[rasn(tag(explicit(context, 0)))]
+    symm_keys: Option<SetOf<rasn::types::OctetString>>,
     #[rasn(tag(explicit(context, 1)))]
     tag1: PCSShareProtectionIdentitiesTag1,
     #[rasn(tag(explicit(context, 2)))]
@@ -567,8 +569,11 @@ pub struct PCSShareProtectionIdentities {
 
 impl PCSShareProtection {
     fn signature_data(&self) -> PCSObjectSignature {
+        // 5 is the version. non-exist is 1, 5 is 2, 4 is 3,
+        // classic is 2
+        // share is 3
         let data = self.attributes.iter().find(|a| a.key == 5).expect("No signature data");
-        rasn::der::decode(&data.value).expect("failed to decode")
+        rasn::der::decode(&data.value).expect("failed to decode signature data")
     }
 
     fn digest_data(&self, objsig: &PCSObjectSignature) -> Vec<u8> {
@@ -577,12 +582,15 @@ impl PCSShareProtection {
             &self.meta[..],
             &objsig.unk2.to_be_bytes(),
             &objsig.unk1.to_be_bytes(),
-            &0u32.to_be_bytes(),
+            &objsig.symm_key_count.unwrap_or(0).to_be_bytes(),
             &objsig.public.keytype.to_be_bytes(),
             &objsig.public.pub_key[..],
         ].concat();
-        if let Some(keylist) = &objsig.keylist {
-            data.extend_from_slice(&rasn::der::encode(keylist).unwrap());
+        if let Some(attributes) = &objsig.attributes {
+            data.extend_from_slice(&rasn::der::encode(attributes).unwrap());
+        }
+        if let Some(ec_key_list) = &objsig.ec_key_list {
+            data.extend_from_slice(&rasn::der::encode(ec_key_list).unwrap());
         }
         data
     }
@@ -599,7 +607,7 @@ impl PCSShareProtection {
         Ok(self.keyset.keyset.first().unwrap().decryption_key.pub_key.to_vec())
     }
 
-    pub fn decrypt_with_keychain(&self, keychain: &KeychainClientState, service: &PCSService<'_>) -> Result<(PCSKey, Vec<CompactECKey<Private>>), PushError> {
+    pub fn decrypt_with_keychain(&self, keychain: &KeychainClientState, service: &PCSService<'_>) -> Result<(Vec<PCSKey>, Vec<CompactECKey<Private>>), PushError> {
         info!("Decoding with {}", base64_encode(&self.decode_key_public()?));
         let account = Value::String(base64_encode(&self.decode_key_public()?));
         let item = keychain.items[service.zone].keys.values().find(|x| x.get("acct").expect("No acct?") == &account).ok_or(PushError::ShareKeyNotFound)?;
@@ -624,6 +632,7 @@ impl PCSShareProtection {
         keyset.make_checksum();
 
         let identities = PCSShareProtectionIdentities {
+            symm_keys: None,
             tag1: Default::default(),
             identities: if keys.is_empty() { None } else { Some(BTreeSet::from_iter([
                 PCSShareProtectionIdentityData {
@@ -666,10 +675,13 @@ impl PCSShareProtection {
                 pub_key: master_ec_key.public_key().to_bytes(master_ec_key.group(), PointConversionForm::UNCOMPRESSED, &mut num_ctx)?.into(),
             },
             signature: Default::default(),
-            keylist: if keys.is_empty() { None } else { Some(keys.iter().map(|k| PCSKeyRef {
+            ec_key_list: if keys.is_empty() { None } else { Some(keys.iter().map(|k| PCSKeyRef {
                 keytype: 3,
                 pub_key: k.compress().to_vec().into(),
-            }).collect()) }
+            }).collect()) },
+            symm_key_count: None,
+            signature_2: None,
+            attributes: None,
         };
 
         let digest_data = protection.digest_data(&signature);
@@ -705,7 +717,8 @@ impl PCSShareProtection {
         Ok(protection)
     }
 
-    pub fn decode(&self, key: &CompactECKey<Private>) -> Result<(PCSKey, Vec<CompactECKey<Private>>), PushError> {
+    pub fn decode(&self, key: &CompactECKey<Private>) -> Result<(Vec<PCSKey>, Vec<CompactECKey<Private>>), PushError> {
+        info!("Decoding share protection!");
         let master_key = PCSKey::new(key, &self.keyset.keyset.first().unwrap().ciphertext)?;
 
         let sig = self.signature_data();
@@ -723,7 +736,15 @@ impl PCSShareProtection {
         let mut verifier = Verifier::new(MessageDigest::sha256(), &key)?;
         verifier.update(&digest_data)?;
         if !verifier.verify(&sig.signature.signature)? {
-            panic!("self sig check failed")
+            if let Some(past_signature) = &sig.signature_2 {
+                let mut verifier = Verifier::new(MessageDigest::sha256(), &key)?;
+                verifier.update(&digest_data)?;
+                if !verifier.verify(&past_signature.signature)? {
+                    panic!("self sig 1 and 2 check failed")
+                }
+            } else {
+                panic!("self sig check failed")
+            }
         }
 
         let hmackey = kdf_ctr_hmac(&master_key.0, "hmackey-of-masterkey".as_bytes(), &[], master_key.0.len());
@@ -747,7 +768,10 @@ impl PCSShareProtection {
             }
         }
 
-        Ok((master_key, keys))
+        let mut pcs_keys = vec![master_key];
+        pcs_keys.extend(identities.symm_keys.unwrap_or_default().into_iter().map(|symm| PCSKey(symm.to_vec())));
+
+        Ok((pcs_keys, keys))
     }
 }
 
@@ -758,6 +782,13 @@ pub struct PCSObjectSignature {
     unk2: u32,
     public: PCSKeyRef,
     signature: PCSSignature,
+    // the ignore fields show up in weird situations, when there are multiple keys?
+    #[rasn(tag(explicit(context, 0)))]
+    symm_key_count: Option<u32>,
+    #[rasn(tag(explicit(context, 1)))]
+    signature_2: Option<PCSSignature>,
     #[rasn(tag(explicit(context, 2)))]
-    keylist: Option<SequenceOf<PCSKeyRef>>,
+    ec_key_list: Option<SequenceOf<PCSKeyRef>>,
+    #[rasn(tag(explicit(context, 3)))]
+    attributes: Option<SequenceOf<PCSAttribute>>,
 }

Original file line number	Diff line number	Diff line change
`@@ -177,4 +177,6 @@ pub enum PushError {`
`177`	`177`	`BatchError(Arc<PushError>),`
`178`	`178`	`#[error("Invalid 2fa code!")]`
`179`	`179`	`Bad2FaCode,`
	`180`	`+ #[error("PCS record key id not found!")]`
	`181`	`+ PCSRecordKeyMissing,`
`180`	`182`	`}`