Fixes

TaeHagen · TaeHagen · commit 08095073413b · 2025-09-04T15:55:05.000-04:00
diff --git a/cloudkit-proto/src/lib.rs b/cloudkit-proto/src/lib.rs
@@ -123,8 +123,7 @@ disjoint_impls! {
 
     impl<T: Serialize + DeserializeOwned + CloudKitBytesKind<Kind = sealed::PlistKind>> CloudKitBytes for T {
         fn from_bytes(v: Vec<u8>) -> Self {
-            // println!("{}", encode_hex(&v));
-            plist::from_bytes(&v).expect("Deserialization failed!")
+            plist::from_bytes(&v).expect(&format!("Deserialization failed! {}", encode_hex(&v)))
         }
 
         fn to_bytes(&self) -> Vec<u8> {
diff --git a/src/cloudkit.rs b/src/cloudkit.rs
@@ -3,8 +3,9 @@ use std::{collections::HashMap, io::{Cursor, Read, Write}, marker::PhantomData,
 use aes::{cipher::consts::U12, Aes128, Aes256};
 use aes_gcm::{AesGcm, Nonce, Tag};
 use aes_siv::siv::CmacSiv;
-use cloudkit_proto::{record, request_operation::header::IsolationLevel, AssetGetResponse, AssetsToDownload, CloudKitRecord, ProtectionInfo, Record, RecordIdentifier, RecordZoneIdentifier, ResponseOperation, Zone};
+use cloudkit_proto::{record, request_operation::header::IsolationLevel, retrieve_changes_response::RecordChange, AssetGetResponse, AssetsToDownload, CloudKitRecord, ProtectionInfo, Record, RecordIdentifier, RecordZoneIdentifier, ResponseOperation, Zone};
 use hkdf::Hkdf;
+use log::info;
 use omnisette::{AnisetteProvider, ArcAnisetteClient};
 use openssl::{bn::{BigNum, BigNumContext}, ec::{EcGroup, EcKey, EcPoint}, hash::MessageDigest, nid::Nid, pkcs5::pbkdf2_hmac, pkey::{HasPublic, PKey, Private}, sha::sha1, sign::{Signer, Verifier}};
 use plist::Value;
@@ -479,6 +480,30 @@ impl FetchRecordChangesOperation {
             include_mergeable_deltas: None,
         })
     }
+
+    pub async fn do_sync(container: &CloudKitOpenContainer<'_, impl AnisetteProvider>, 
+        zones: &[(cloudkit_proto::RecordZoneIdentifier, Option<Vec<u8>>)], assets: &cloudkit_proto::AssetsToDownload) -> Result<Vec<(Vec<AssetGetResponse>, Vec<RecordChange>, Option<Vec<u8>>)>, PushError> {
+        let mut responses = zones.iter().map(|zone| (vec![], vec![], zone.1.clone())).collect::<Vec<_>>();
+
+        let mut finished_zones = vec![];
+        while finished_zones.len() != zones.len() {
+            let mut sync_zones_here = zones.iter().enumerate().filter(|(_, zone)| !finished_zones.contains(&zone.0)).collect::<Vec<_>>();
+            let operations = container.perform_operations_checked(&CloudKitSession::new(), 
+                &sync_zones_here.iter().map(|(idx, zone)| FetchRecordChangesOperation::new(zone.0.clone(), responses[*idx].2.clone(), assets))
+                                .collect::<Vec<_>>(), IsolationLevel::Zone).await?;
+            for (result, (zone_idx, zone)) in operations.into_iter().zip(sync_zones_here.iter_mut()) {
+                if result.1.status() == 3 {
+                    // done syncing
+                    finished_zones.push(zone.0.clone());
+                }
+                responses[*zone_idx].0.extend(result.0);
+                responses[*zone_idx].1.extend(result.1.change);
+                responses[*zone_idx].2 = result.1.sync_continuation_token.clone();
+            }
+        }
+
+        Ok(responses)
+    }
 }
 
 pub fn should_reset(error: Option<&PushError>) -> bool {
@@ -846,10 +871,13 @@ impl<'t, T: AnisetteProvider> CloudKitOpenContainer<'t, T> {
                 ..
             })) => {
                 let service = PCSPrivateKey::get_service_key(client, pcs_service, self.client.config.as_ref()).await?;
+                
+                info!("Creating zone {} with service key {}", zone_name, encode_hex(&service.key().compress()));
 
-                let request = ZoneSaveOperation::new(zone.clone(), Some(&service.key()), pcs_service.global_record).unwrap();
+                let request = ZoneSaveOperation::new(zone.clone(), Some(&service.key()), pcs_service.global_record)?;
                 let zone = request.0.clone().zone.unwrap();
-                self.perform(&CloudKitSession::new(), request).await.unwrap();
+                self.perform(&CloudKitSession::new(), request).await?;
+                info!("Created zone");
                 zone
             },
             Err(err) => return Err(err)
@@ -978,6 +1006,9 @@ impl<'t, T: AnisetteProvider> CloudKitOpenContainer<'t, T> {
         if response.status().as_u16() == 401 {
             self.client.token_provider.refresh_mme().await?;
         }
+        if response.status().as_u16() == 429 {
+            return Err(PushError::TooManyRequests);
+        }
         
         let token: Vec<u8> = response.bytes().await?
             .into();
diff --git a/src/error.rs b/src/error.rs
@@ -171,8 +171,8 @@ pub enum PushError {
     NotInClique,
     #[error("Missing group photo!")]
     MissingGroupPhoto,
-    #[error("PCS Share key not found!")]
-    ShareKeyNotFound,
+    #[error("PCS Share key {0} not found!")]
+    ShareKeyNotFound(String),
     #[error("BatchError {0}")]
     BatchError(Arc<PushError>),
     #[error("Invalid 2fa code!")]
@@ -181,4 +181,6 @@ pub enum PushError {
     PCSRecordKeyMissing,
     #[error("Circle is over!")]
     CircleOver,
+    #[error("Too many requests!")]
+    TooManyRequests,
 }
diff --git a/src/findmy.rs b/src/findmy.rs
@@ -420,22 +420,20 @@ impl<P: AnisetteProvider> FindMyClient<P> {
 
         let mut state = self.state.state.lock().await;
 
-        let mut result = container.perform(&CloudKitSession::new(),
-            FetchRecordChangesOperation::new(beacon_zone.clone(), state.state_token.clone(), &NO_ASSETS)).await;
+        let mut result = FetchRecordChangesOperation::do_sync(&container, &[(beacon_zone.clone(), state.state_token.clone())], &NO_ASSETS).await;
         if should_reset(result.as_ref().err()) {
             state.state_token = None;
             state.accessories.clear();
-            result = container.perform(&CloudKitSession::new(),
-                FetchRecordChangesOperation::new(beacon_zone, state.state_token.clone(), &NO_ASSETS)).await;
+            result = FetchRecordChangesOperation::do_sync(&container, &[(beacon_zone.clone(), state.state_token.clone())], &NO_ASSETS).await;
         }
 
-        let (_, changes) = result?;
+        let (_, changes, continuation) = result?.remove(0);
         
-        state.state_token = changes.sync_continuation_token.clone();
+        state.state_token = continuation.clone();
 
         let accessories = &mut state.accessories;
         
-        for change in changes.change {
+        for change in changes {
             let identifier = change.identifier.as_ref().unwrap().value.as_ref().unwrap().name().to_string();
             let Some(record) = change.record else {
                 accessories.remove(&identifier);
@@ -447,7 +445,7 @@ impl<P: AnisetteProvider> FindMyClient<P> {
             if record.r#type.as_ref().unwrap().name() == MasterBeaconRecord::record_type() {
                 let item = MasterBeaconRecord::from_record_encrypted(&record.record_field, Some((&pcs_key_for_record(&record, &key)?, record.record_identifier.as_ref().unwrap())));
 
-                println!("Got beacon {:?} {}", item, identifier);
+                info!("Got beacon {:?} {}", item, identifier);
 
                 if let Some(accessory) = accessories.get_mut(&identifier) {
                     accessory.master_record = item;
@@ -524,6 +522,7 @@ impl<P: AnisetteProvider> FindMyClient<P> {
                 let adv = base64_encode(&sha256(&x.to_vec_padded(28)?));
                 key_map.insert(adv.clone(), (id.clone(), key, idx));
                 device_keys.push(adv);
+            
             }
             
             search.push(json!({
@@ -536,6 +535,11 @@ impl<P: AnisetteProvider> FindMyClient<P> {
             }));
         }
 
+        if search.is_empty() {
+            info!("Not searching, no item!");
+            return Ok(())
+        }
+
         let mut request = self.anisette.lock().await.get_headers().await?.clone();
         request.remove("X-Mme-Client-Info").unwrap();
         let anisette_headers: HeaderMap = request.into_iter().map(|(a, b)| (HeaderName::from_str(&a).unwrap(), b.parse().unwrap())).collect();
diff --git a/src/imessage/cloud_messages.proto b/src/imessage/cloud_messages.proto
@@ -13,8 +13,8 @@ message MessageProto2 {
 message MessageProto {
     uint32 unk1 = 1; // always 1, is finished or ck sync state
     optional string groupTitle = 2;
-    string text = 3;
-    bytes attributedBody = 4;
+    optional string text = 3;
+    optional bytes attributedBody = 4;
     optional string balloonBundleId = 5;
     optional bytes payloadData = 6;
     optional bytes messageSummaryInfo = 7;
diff --git a/src/imessage/cloud_messages.rs b/src/imessage/cloud_messages.rs
@@ -6,6 +6,7 @@ use std::ops::{Deref, DerefMut};
 use std::sync::Arc;
 use std::time::{Duration, SystemTime};
 
+use backon::{ConstantBuilder, Retryable};
 use cloudkit_derive::CloudKitRecord;
 use cloudkit_proto::request_operation::header::IsolationLevel;
 use cloudkit_proto::retrieve_changes_response::RecordChange;
@@ -400,7 +401,7 @@ impl Into<Option<MMCSAttachmentMeta>> for &Attachment {
 #[derive(Serialize, Deserialize, Debug, Clone, Default)]
 pub struct AttachmentMetaExtra {
     #[serde(rename = "pgens")]
-    pub preview_generation_state: Option<u32>, // set to 1
+    pub preview_generation_state: Option<i32>, // set to 1
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone, Default)]
@@ -412,7 +413,7 @@ pub struct AttachmentMeta {
     #[serde(rename = "tb")]
     pub total_bytes: u64,
     #[serde(rename = "st")]
-    pub transfer_state: u32,
+    pub transfer_state: i32,
     #[serde(rename = "is")]
     pub is_sticker: bool,
     #[serde(rename = "aguid")]
@@ -430,7 +431,7 @@ pub struct AttachmentMeta {
     #[serde(rename = "tn")]
     pub transfer_name: String,
     #[serde(rename = "vers")]
-    pub version: u32, // set to 1
+    pub version: i32, // set to 1
     #[serde(rename = "t")]
     pub uti: Option<String>, // uti type
     #[serde(rename = "cdt")]
@@ -474,26 +475,17 @@ impl<P: AnisetteProvider> CloudMessagesClient<P> {
         return Ok(locked.clone().unwrap())
     }
 
-    // todo keep syncing until end
-    async fn sync_records<T: CloudKitRecord>(&self, zone: &str, mut continuation_token: Option<Vec<u8>>) -> Result<(Vec<u8>, HashMap<String, Option<T>>), PushError> {
+    async fn sync_records<T: CloudKitRecord>(&self, zone: &str, continuation_token: Option<Vec<u8>>) -> Result<(Vec<u8>, HashMap<String, Option<T>>, i32), PushError> {
         let container = self.get_container().await?;
 
         let zone = container.private_zone(zone.to_string());
         let key = container.get_zone_encryption_config(&zone, &self.keychain, &MESSAGES_SERVICE).await?;
-        let (_assets, mut response) = container.perform(&CloudKitSession::new(),
+        let (_assets, response) = container.perform(&CloudKitSession::new(),
             FetchRecordChangesOperation::new(zone.clone(), continuation_token, &NO_ASSETS)).await?;
-        let mut items = vec![];
-        while !response.change.is_empty() {
-            items.extend(response.change);
-            continuation_token = response.sync_continuation_token.clone();
-            response = container.perform(&CloudKitSession::new(),
-                FetchRecordChangesOperation::new(zone.clone(), continuation_token, &NO_ASSETS)).await?.1;
-        }
-
 
         let mut results = HashMap::new();
 
-        for change in &items {
+        for change in &response.change {
             let identifier = change.identifier.as_ref().unwrap().value.as_ref().unwrap().name().to_string();
 
             let Some(record) = &change.record else {
@@ -515,7 +507,7 @@ impl<P: AnisetteProvider> CloudMessagesClient<P> {
             results.insert(identifier, Some(item));
         }
 
-        Ok((response.sync_continuation_token().to_vec(), results))
+        Ok((response.sync_continuation_token().to_vec(), results, response.status()))
     }
 
     async fn save_records<T: CloudKitRecord>(&self, zone: &str, records: HashMap<String, T>) -> Result<HashMap<String, Result<(), PushError>>, PushError> {
@@ -555,13 +547,16 @@ impl<P: AnisetteProvider> CloudMessagesClient<P> {
 
         let zone = container.private_zone(zone.to_string());
 
-        let mut operations = vec![];
-        for record_id in records {
-            operations.push(DeleteRecordOperation::new(record_identifier(zone.clone(), record_id)));
+        for batch in records.chunks(256) {
+            let mut operations = vec![];
+            for record_id in batch {
+                operations.push(DeleteRecordOperation::new(record_identifier(zone.clone(), record_id)));
+            }
+            (|| async {
+                container.perform_operations_checked(&CloudKitSession::new(), &operations, IsolationLevel::Operation).await
+            }).retry(&ConstantBuilder::default().with_delay(Duration::from_secs(5)).with_max_times(3)).await?;
         }
 
-        container.perform_operations_checked(&CloudKitSession::new(), &operations, IsolationLevel::Operation).await?;
-
         Ok(())
     }
 
@@ -575,11 +570,18 @@ impl<P: AnisetteProvider> CloudMessagesClient<P> {
             ZoneDeleteOperation::new(container.private_zone("messageManateeZone".to_string())),
             ZoneDeleteOperation::new(container.private_zone("attachmentManateeZone".to_string())),
             ZoneDeleteOperation::new(container.private_zone("chat1ManateeZone".to_string())),
+            ZoneDeleteOperation::new(container.private_zone("messageUpdateZone".to_string())),
+            ZoneDeleteOperation::new(container.private_zone("recoverableMessageDeleteZone".to_string())),
+            ZoneDeleteOperation::new(container.private_zone("scheduledMessageZone".to_string())),
+            ZoneDeleteOperation::new(container.private_zone("chatBotMessageZone".to_string())),
+            ZoneDeleteOperation::new(container.private_zone("chatBotAttachmentZone".to_string())),
+            ZoneDeleteOperation::new(container.private_zone("chatBotRecoverableMessageDeleteZone".to_string())),
+
         ], IsolationLevel::Operation).await?;
         Ok(())
     }
 
-    pub async fn sync_chats(&self, continuation_token: Option<Vec<u8>>) -> Result<(Vec<u8>, HashMap<String, Option<CloudChat>>), PushError> {
+    pub async fn sync_chats(&self, continuation_token: Option<Vec<u8>>) -> Result<(Vec<u8>, HashMap<String, Option<CloudChat>>, i32), PushError> {
         self.sync_records("chatManateeZone", continuation_token).await
     }
 
@@ -591,7 +593,7 @@ impl<P: AnisetteProvider> CloudMessagesClient<P> {
         self.delete_records("chatManateeZone", chats).await
     }
 
-    pub async fn sync_messages(&self, continuation_token: Option<Vec<u8>>) -> Result<(Vec<u8>, HashMap<String, Option<CloudMessage>>), PushError> {
+    pub async fn sync_messages(&self, continuation_token: Option<Vec<u8>>) -> Result<(Vec<u8>, HashMap<String, Option<CloudMessage>>, i32), PushError> {
         self.sync_records("messageManateeZone", continuation_token).await
     }
 
@@ -603,7 +605,7 @@ impl<P: AnisetteProvider> CloudMessagesClient<P> {
         self.delete_records("messageManateeZone", messages).await
     }
 
-    pub async fn sync_attachments(&self, continuation_token: Option<Vec<u8>>) -> Result<(Vec<u8>, HashMap<String, Option<CloudAttachment>>), PushError> {
+    pub async fn sync_attachments(&self, continuation_token: Option<Vec<u8>>) -> Result<(Vec<u8>, HashMap<String, Option<CloudAttachment>>, i32), PushError> {
         self.sync_records("attachmentManateeZone", continuation_token).await
     }
 
diff --git a/src/keychain.rs b/src/keychain.rs
@@ -1,4 +1,4 @@
-use std::{collections::{BTreeMap, HashMap}, io::{Cursor, Read}, sync::Arc};
+use std::{collections::{BTreeMap, HashMap}, io::{Cursor, Read}, sync::Arc, time::Duration};
 
 use aes_gcm::{AesGcm, Nonce};
 use cloudkit_derive::CloudKitRecord;
@@ -32,6 +32,8 @@ use srp::{client::{SrpClient, SrpClientVerifier}, groups::G_2048, server::SrpSer
 use tokio::sync::{Mutex, RwLock};
 use crate::{aps::APSInterestToken, auth::MobileMeDelegateResponse, cloudkit::{CloudKitClient, CloudKitContainer, CloudKitOpenContainer, CloudKitSession, FetchRecordChangesOperation, FunctionInvokeOperation, ALL_ASSETS}, ids::{CompactECKey, IDSRecvMessage}, util::{base64_decode, base64_encode, bin_deserialize, bin_deserialize_opt_vec, bin_serialize, bin_serialize_opt_vec, decode_hex, decode_uleb128, duration_since_epoch, ec_deserialize_priv, ec_serialize_priv, encode_hex, kdf_ctr_hmac, plist_to_bin, plist_to_string, proto_deserialize, proto_deserialize_opt, proto_serialize, proto_serialize_opt, rfc6637_unwrap_key, NSData, NSDataClass, REQWEST}, APSConnection, APSMessage, IdentityManager, KeyedArchive, OSConfig, PushError};
 
+use backon::{BackoffBuilder, ConstantBuilder, ExponentialBuilder};
+use backon::Retryable;
 
 // ansi_x963_kdf (use crate later, current dependencies are broken)
 #[inline]
@@ -1161,23 +1163,21 @@ impl<P: AnisetteProvider> KeychainClient<P> {
         let security_container = self.get_security_container().await?;
 
         let mut state = self.state.write().await;
-        let mut result = security_container.perform_operations_checked(&CloudKitSession::new(), 
-            &zones.iter().map(|zone| FetchRecordChangesOperation::new(security_container.private_zone(zone.to_string()), 
-                state.items.get(*zone).and_then(|z| z.change_tag.clone()).map(|z| z.into()), &ALL_ASSETS)).collect::<Vec<_>>(), IsolationLevel::Zone).await;
+        let mut result = FetchRecordChangesOperation::do_sync(&security_container, &zones.iter()
+            .map(|zone| (security_container.private_zone(zone.to_string()), state.items.get(*zone).and_then(|z| z.change_tag.clone()).map(|z| z.into()))).collect::<Vec<_>>(), &ALL_ASSETS).await;
         if should_reset(result.as_ref().err()) {
             state.items.clear();
-            result = security_container.perform_operations_checked(&CloudKitSession::new(), 
-                &zones.iter().map(|zone| FetchRecordChangesOperation::new(security_container.private_zone(zone.to_string()), 
-                    state.items.get(*zone).and_then(|z| z.change_tag.clone()).map(|z| z.into()), &ALL_ASSETS)).collect::<Vec<_>>(), IsolationLevel::Zone).await;
+            result = FetchRecordChangesOperation::do_sync(&security_container, &zones.iter()
+            .map(|zone| (security_container.private_zone(zone.to_string()), state.items.get(*zone).and_then(|z| z.change_tag.clone()).map(|z| z.into()))).collect::<Vec<_>>(), &ALL_ASSETS).await;
         }
         let item = result?;
 
         let state = &mut *state;
 
-        for (zone, (_, item)) in zones.iter().zip(item.into_iter()) {
+        for (zone, (_, changes, change)) in zones.iter().zip(item.into_iter()) {
             let saved_keychain_zone = state.items.entry(zone.to_string()).or_default();
-            saved_keychain_zone.change_tag = Some(item.sync_continuation_token.unwrap().into());
-            for change in item.change {
+            saved_keychain_zone.change_tag = change.clone().map(|i| i.into());
+            for change in changes {
                 let Some(record) = change.record else {
                     warn!("record missing change {:?}", change);
                     continue
@@ -1309,13 +1309,29 @@ impl<P: AnisetteProvider> KeychainClient<P> {
 
         drop(state);
 
-        security.perform_operations_checked(&CloudKitSession::new(), &delete_ops, IsolationLevel::Zone).await?;
 
-        let _: CuttlefishResetResponse = self.invoke_cuttlefish("reset", CuttlefishResetRequest {
-            reason: Some(3),
-        }).await?;
+        (|| async {
+            (|| async { security.perform_operations_checked(&CloudKitSession::new(), &delete_ops, IsolationLevel::Zone).await })
+                .retry(&ConstantBuilder::default()
+                .with_delay(Duration::ZERO)
+                .with_max_times(3))
+                .notify(|err: &PushError, _dur: Duration| {
+                    println!("erasing keys {:?}", err);
+            }).await?;
+
+            let _: CuttlefishResetResponse = self.invoke_cuttlefish("reset", CuttlefishResetRequest {
+                reason: Some(3),
+            }).await?;
 
-        self.join_clique(device_password, None, &shares, viewkeys).await?;
+            self.join_clique(device_password, None, &shares, viewkeys.clone()).await?;
+            Ok(())
+        })
+            .retry(&ConstantBuilder::default()
+            .with_delay(Duration::ZERO)
+            .with_max_times(2))
+            .notify(|err: &PushError, _dur: Duration| {
+                println!("resetting clique {:?}", err);
+            }).await?;
 
         Ok(())
     }
diff --git a/src/pcs.rs b/src/pcs.rs
diff --git a/src/util.rs b/src/util.rs

Original file line number	Diff line number	Diff line change
`@@ -123,8 +123,7 @@ disjoint_impls! {`
`123`	`123`
`124`	`124`	`impl<T: Serialize + DeserializeOwned + CloudKitBytesKind<Kind = sealed::PlistKind>> CloudKitBytes for T {`
`125`	`125`	`fn from_bytes(v: Vec<u8>) -> Self {`
`126`		`- // println!("{}", encode_hex(&v));`
`127`		`- plist::from_bytes(&v).expect("Deserialization failed!")`
	`126`	`+ plist::from_bytes(&v).expect(&format!("Deserialization failed! {}", encode_hex(&v)))`
`128`	`127`	`}`
`129`	`128`
`130`	`129`	`fn to_bytes(&self) -> Vec<u8> {`