Fixes

TaeHagen · TaeHagen · commit b2996d57936c · 2026-04-26T21:05:10.000-07:00
diff --git a/src/aps.rs b/src/aps.rs
@@ -237,7 +237,8 @@ impl APSPackedEncoder {
         } else {
             // search in tag cache
             if let Some(existing) = self.key_table.get_buf_idx(&attribute.id) {
-                self.write_tagged_value(5, existing as u64, 0x20 | flags, writer)?;
+                info!("Using kcache {existing}");
+                writer.write_all(&[0x20 | flags | existing as u8])?;
             } else {
                 self.write_tagged_value(4, attribute.id.len() as u64 - 1, 0x10 | flags, writer)?;
 
@@ -248,6 +249,7 @@ impl APSPackedEncoder {
 
         // actually write the value
         if let Some(existing) = existing {
+            info!("Using vcache {existing}");
             self.write_tagged_value(8, existing as u64, 0, writer)?;
         } else if let Some(cached_buf) = cached_buf {
             writer.write_all(&cached_buf)?;
@@ -945,6 +947,148 @@ impl APSMessage {
 }
 
 
+#[tokio::test]
+async fn replay_test() {
+    use keystore::{init_keystore, software::{SoftwareKeystore, NoEncryptor}};
+    init_keystore(SoftwareKeystore {
+        state: plist::from_file("jerrytest/keystore.plist").unwrap(),
+        update_state: Box::new(|state| {
+            plist::to_file_xml("jerrytest/keystore.plist", state).unwrap();
+        }),
+        encryptor: NoEncryptor,
+    });
+
+    if let Err(_) = std::env::var("RUST_LOG") {
+        std::env::set_var("RUST_LOG", "debug");
+    }
+    let _ = pretty_env_logger::try_init();
+
+    let state_path = std::env::var("APS_REPLAY_STATE")
+        .unwrap_or_else(|_| if std::path::Path::new("jerrytest/push.plist").exists() {
+            "jerrytest/push.plist".to_string()
+        } else {
+            "jerrytest/push.plist".to_string()
+        });
+    let log_path = std::env::var("APS_REPLAY_LOG")
+        .unwrap_or_else(|_| if std::path::Path::new("jerrytest/replaytest.log").exists() {
+            "jerrytest/replaytest.log".to_string()
+        } else {
+            "jerrytest/replaytest.log".to_string()
+        });
+
+    let mut state: APSState = plist::from_file(&state_path).unwrap();
+    let pair = state.keypair.as_ref().expect("replay APS state must already contain a keypair");
+
+    let (socket, encoder, mut decoder) = open_socket().await.unwrap();
+    let (mut read, write) = split(socket);
+    let (send, _) = tokio::sync::broadcast::channel(999);
+    let socket = Arc::new(DebugMutex::new(Some((write, encoder))));
+
+    let reader_send = send.clone();
+    let mut reader = task::spawn(async move {
+        loop {
+            match APSMessage::read_from_stream(&mut read, &mut decoder).await {
+                Ok(Some(msg)) => {
+                    let _ = reader_send.send(msg);
+                },
+                Ok(None) => {},
+                Err(err) => {
+                    return Err(err);
+                }
+            };
+        }
+
+        #[allow(unreachable_code)]
+        Ok::<(), PushError>(())
+    });
+
+    let nonce = generate_nonce(NonceType::APNS);
+    let signature = do_ids_signature(&pair.private, &nonce).unwrap();
+    let mut recv = send.subscribe();
+    replay_send(&socket, APSMessage::Connect {
+        flags: 0b01000001,
+        certificate: Some(pair.cert.clone()),
+        nonce: Some(nonce),
+        signature: Some(signature),
+        token: state.token,
+    }).await.unwrap();
+
+    let (token, status) = tokio::time::timeout(Duration::from_secs(15), async {
+        loop {
+            if let APSMessage::ConnectResponse { token, status } = recv.recv().await.unwrap() {
+                return (token, status);
+            }
+        }
+    }).await.unwrap();
+
+    assert_eq!(status, 0, "APS connect failed with status {status}");
+    if let Some(token) = token {
+        state.token = Some(token);
+    }
+
+    let replay = std::fs::read_to_string(&log_path).unwrap();
+    let mut sent = 0usize;
+    let mut skipped_first_send = false;
+    for (line_index, line) in replay.lines().enumerate() {
+        let Some((_, rest)) = line.split_once("Sending \"") else {
+            continue;
+        };
+        if !skipped_first_send {
+            skipped_first_send = true;
+            continue;
+        }
+        if sent >= 0 {
+            break;
+        }
+        let Some((hex, _)) = rest.split_once('"') else {
+            panic!("bad Sending line {}: missing closing quote", line_index + 1);
+        };
+        info!("Sending {}", sent + 1);
+        let bytes = decode_hex(hex).unwrap();
+        let message: APSMessage = plist::from_bytes(&bytes).unwrap();
+        replay_send(&socket, message).await.unwrap();
+        sent += 1;
+    }
+
+    assert!(skipped_first_send, "no Sending lines found in {log_path}");
+    info!("replayed {sent} APS messages from {log_path}");
+
+    tokio::select! {
+        result = &mut reader => {
+            match result {
+                Ok(Ok(())) => panic!("replay APS reader ended before observation window ended"),
+                Ok(Err(err)) => panic!("replay APS reader stopped before observation window ended: {err}"),
+                Err(err) => panic!("replay APS reader task failed before observation window ended: {err}"),
+            }
+        }
+        _ = tokio::time::sleep(Duration::from_secs(5)) => {}
+    }
+
+    reader.abort();
+}
+
+#[cfg(test)]
+async fn replay_send(
+    socket: &DebugMutex<Option<(WriteHalf<TlsStream<TcpStream>>, Option<APSPackedEncoder>)>>,
+    message: APSMessage,
+) -> Result<(), PushError> {
+    let mut socket_guard = socket.lock().await;
+    let socket = socket_guard.as_mut().ok_or(PushError::NotConnected)?;
+    if let Some(encoder) = &mut socket.1 {
+        let mut buf = vec![];
+        encoder.encode_message(message.to_packed_raw(), Cursor::new(&mut buf))?;
+        socket.0.write_all(&buf).await?;
+    } else {
+        let mut raw = message.to_raw();
+        for message in &mut raw.body {
+            message.update()?;
+        }
+        raw.update()?;
+        socket.0.write_all(&raw.to_bytes()?).await?;
+    }
+    Ok(())
+}
+
 #[tokio::test]
 async fn proxy() {
     if let Err(_) = std::env::var("RUST_LOG") {
@@ -1092,11 +1236,16 @@ async fn open_socket() -> Result<(TlsStream<TcpStream>, Option<APSPackedEncoder>
     if let Some(protocol) = stream.get_ref().1.alpn_protocol() {
         let protocol = std::str::from_utf8(protocol).unwrap();
         if protocol.starts_with("apns-pack-v1") {
-            let mut parts = protocol.split(":");
-            let encoder_count: usize = parts.nth(1).expect("Bad ALPN1!").parse().expect("Bad alpn3!");
-            let decoder_count: usize = parts.next().expect("Bad ALPN2!").parse().expect("Bad alpn4!");
-            encoder = Some(APSPackedEncoder::new_with_cache_size(encoder_count));
-            decoder = Some(APSPackedDecoder::new_with_cache_size(decoder_count));
+            if protocol.contains(":") {
+                let mut parts = protocol.split(":");
+                let encoder_count: usize = parts.nth(1).expect("Bad ALPN1!").parse().expect("Bad alpn3!");
+                let decoder_count: usize = parts.next().expect("Bad ALPN2!").parse().expect("Bad alpn4!");
+                encoder = Some(APSPackedEncoder::new_with_cache_size(encoder_count));
+                decoder = Some(APSPackedDecoder::new_with_cache_size(decoder_count));
+            } else {
+                encoder = Some(APSPackedEncoder::default());
+                decoder = Some(APSPackedDecoder::default());
+            }
         }
     }
 
@@ -1387,7 +1536,8 @@ impl APSConnectionResource {
 
             let mut buf = vec![];
             encoder.encode_message(message.to_packed_raw(), Cursor::new(&mut buf))?;
-
+            
+            debug!("Sendin2g {:?}", encode_hex(&buf));
             socket.0.write_all(&buf).await
         } else {
             let mut raw = message.to_raw();
diff --git a/src/statuskit.rs b/src/statuskit.rs
@@ -716,59 +716,60 @@ impl<T: AnisetteProvider + Send + Sync + 'static> StatusKitClient<T> {
     }
 
     pub async fn handle(&self, msg: APSMessage) -> Result<Option<StatusKitMessage>, PushError> {
-        let APSMessage::Notification { id: _, topic, token: _, payload: Value::Data(payload), channel } = msg.clone() else { return Ok(None) };
-        if let Some(channel) = &channel {
-            let mut state = self.state.write().await;
-            if let Some(local_channel) = state.recent_channels.iter_mut().find(|c| c.identifier.id == channel.id && sha1(c.identifier.topic.as_bytes()) == topic) {
-                local_channel.last_msg_ns = channel.last_msg_ns;
-                (self.update_state)(&state);
+        if let APSMessage::Notification { id: _, topic, token: _, payload: Value::Data(payload), channel } = msg.clone() {
+            if let Some(channel) = &channel {
+                let mut state = self.state.write().await;
+                if let Some(local_channel) = state.recent_channels.iter_mut().find(|c| c.identifier.id == channel.id && sha1(c.identifier.topic.as_bytes()) == topic) {
+                    local_channel.last_msg_ns = channel.last_msg_ns;
+                    (self.update_state)(&state);
+                }
             }
-        }
 
-        if topic == sha1("com.apple.icloud.presence.mode.status".as_bytes()) {
-            let result: StatusKitOuterMessage = serde_json::from_slice(&payload)?;
-            let inner: StatusKitInnerMessage = plist::from_bytes(&base64_decode(&result.status_kit_data_key))?;
+            if topic == sha1("com.apple.icloud.presence.mode.status".as_bytes()) {
+                let result: StatusKitOuterMessage = serde_json::from_slice(&payload)?;
+                let inner: StatusKitInnerMessage = plist::from_bytes(&base64_decode(&result.status_kit_data_key))?;
 
-            
-            let Some(channel) = &channel else { return Ok(None) };
-            debug!("Got statuskit message {inner:?} on channel {}", encode_hex(&channel.id));
-            let mut state = self.state.write().await;
-            let Some(referenced_channel) = state.keys.get_mut(&base64_encode(&channel.id)) else { panic!("Channel not found!") };
+                
+                let Some(channel) = &channel else { return Ok(None) };
+                debug!("Got statuskit message {inner:?} on channel {}", encode_hex(&channel.id));
+                let mut state = self.state.write().await;
+                let Some(referenced_channel) = state.keys.get_mut(&base64_encode(&channel.id)) else { panic!("Channel not found!") };
 
-            let key = referenced_channel.get_key(inner.ratchet)?;
+                let key = referenced_channel.get_key(inner.ratchet)?;
 
-            let message = base64_decode(&inner.encrypted_message);
-            referenced_channel.signature.verify(MessageDigest::sha256(), &message, base64_decode(&inner.signature).try_into().expect("Bad signature length!"))?;
+                let message = base64_decode(&inner.encrypted_message);
+                referenced_channel.signature.verify(MessageDigest::sha256(), &message, base64_decode(&inner.signature).try_into().expect("Bad signature length!"))?;
 
-            let key = key.message_key();
+                let key = key.message_key();
 
-            let cipher = Aes256Gcm::new_from_slice(&key).expect("GCM key creation failed");
-            let plaintext = cipher.decrypt(Nonce::from_slice(&message[..12]), aes_gcm::aead::Payload {
-                msg: &message[12..],
-                aad: &referenced_channel.signature.compress()
-            }).expect("Failed to decrypt");
+                let cipher = Aes256Gcm::new_from_slice(&key).expect("GCM key creation failed");
+                let plaintext = cipher.decrypt(Nonce::from_slice(&message[..12]), aes_gcm::aead::Payload {
+                    msg: &message[12..],
+                    aad: &referenced_channel.signature.compress()
+                }).expect("Failed to decrypt");
 
-            let status = PublishedStatus::decode(Cursor::new(&plaintext))?;
-            let status: StatusKitStatus = plist::from_bytes(&status.message)?;
+                let status = PublishedStatus::decode(Cursor::new(&plaintext))?;
+                let status: StatusKitStatus = plist::from_bytes(&status.message)?;
 
-            let user = referenced_channel.from.clone();
-            let is_available = status.active || status.id.as_ref().map(|s| referenced_channel.personal_config.allowed_modes.contains(s)).unwrap_or(false);
+                let user = referenced_channel.from.clone();
+                let is_available = status.active || status.id.as_ref().map(|s| referenced_channel.personal_config.allowed_modes.contains(s)).unwrap_or(false);
 
-            (self.update_state)(&state); // we might have ratcheted it
+                (self.update_state)(&state); // we might have ratcheted it
 
 
-            return Ok(Some(StatusKitMessage::StatusChanged {
-                user,
-                mode: status.id.clone(),
-                allowed: is_available
-            }));
+                return Ok(Some(StatusKitMessage::StatusChanged {
+                    user,
+                    mode: status.id.clone(),
+                    allowed: is_available
+                }));
+            }
         }
 
         if let Some(IDSRecvMessage { message_unenc: Some(message), topic, sender: Some(sender), .. }) = self.identity.receive_message(msg, &["com.apple.private.alloy.status.keysharing", "com.apple.private.alloy.status.personal"]).await? {
             match topic {
                 "com.apple.private.alloy.status.keysharing" => {
                     let parsed: StatusKitRawSharedDevice = message.plist()?;
-                    debug!("StatusKit IDS message came in as {:?}", parsed);
+                    info!("StatusKit IDS message came in as {:?}", parsed);
 
                     let config: StatusKitPersonalConfig = plist::from_bytes(&base64_decode(&parsed.personal_config))?;
                     let share_message = SharedMessage::decode(Cursor::new(base64_decode(&parsed.keys)))?;
diff --git a/src/test.rs b/src/test.rs
@@ -404,7 +404,7 @@ async fn main() {
         let pet = account.get_pet().unwrap();
         let spd = account.spd.as_ref().unwrap();
 
-        let delegates = login_apple_delegates(&gsa.user, &pet, spd["adsid"].as_string().unwrap(), None, &mut *anisette_client.lock().await, config.as_ref(), &[LoginDelegate::IDS, LoginDelegate::MobileMe]).await.unwrap();
+        let delegates = login_apple_delegates(&account, None, config.as_ref(), &[LoginDelegate::IDS, LoginDelegate::MobileMe]).await.unwrap();
         let user = authenticate_apple(delegates.ids.unwrap(), config.as_ref()).await.unwrap();
 
         let mobileme = delegates.mobileme.unwrap();
@@ -605,10 +605,10 @@ async fn main() {
     token_provider.clone(), anisette_client.clone(), client.identity.clone()).await.unwrap();
 
     let state: PasswordState = plist::from_file("passwords.plist").unwrap_or_default();
-    let passwords = PasswordManager::new(
-        keychain.clone(), cloudkit.clone(), client.identity.clone(), connection.clone(), state, Box::new(move |state| {
-            plist::to_file_xml("passwords.plist", state).unwrap();
-        })).await;
+    // let passwords = PasswordManager::new(
+    //     keychain.clone(), cloudkit.clone(), client.identity.clone(), connection.clone(), state, Box::new(move |state| {
+    //         plist::to_file_xml("passwords.plist", state).unwrap();
+    //     })).await;
 
 
     if let Some(mut s) = session {
@@ -933,9 +933,9 @@ async fn main() {
         tokio::select! {
             msg = subscription.recv() => {
                 let msg = msg.unwrap();
-                if let Err(e) = passwords.handle(msg.clone()).await {
-                    info!("err {e}");
-                }
+                // if let Err(e) = passwords.handle(msg.clone()).await {
+                //     info!("err {e}");
+                // }
                 // if let Err(e) = findmy_client.handle(msg.clone()).await {
                 //     info!("err {e}");
                 // }
