Fix reconnection with a cached base token

[akinazuki]

Summary

• Fix _receive hanging forever when reconnecting with a pre-cached base_token, caused by the ConnectAck filter requiring an exact token match
• Fix crash when server accepts the cached token without returning one in ConnectAck (ack.token is None)
• Fix race condition where base_token property returned before the TCP connection was established

Background

When the token parameter is passed to create_apns_connection (e.g. restoring a previously persisted token), three things go wrong:

1. The ConnectAck receive filter checks c.token == self._base_token. But the server may return a different token, or omit the token field entirely (item 0x03 absent) when it accepts the cached one. This causes _receive to wait forever.

2. The base_token property skips _connected.wait() when _base_token is already set, so callers can attempt to send commands before _conn is initialized.

3. assert ack.token == self._base_token crashes when ack.token is None.

This matches real apsd behavior: the daemon persists the base token to com.apple.apsd.plist and sends it on reconnect (e.g. after a device reboot). The server typically accepts it silently.

[JJTech0130]

But the server may return a different token, or omit the token field entirely (item 0x03 absent) when it accepts the cached one.

Have you actually observed the server behaving in this manner? I don't recall encountering this during testing (I must admit it was a long time ago though).

The base_token property skips _connected.wait() when _base_token is already set, so callers can attempt to send commands before _conn is initialized.

It appears that _update_filter and mint_scoped_token are the only functions which rely on this behavior, it's possible they might call _send before the connected event fires.
_conn is probably still None in that case in which case it will bail with an assert... yeah, that looks like an issue.

If you can send a log proving that you encountered the above issue and it wasn't just an AI hallucinating it, then I can probably accept this PR.

[akinazuki]

here is an poc

import asyncio
import json
import os
import time
import logging

from cryptography import x509
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.backends import default_backend

logging.basicConfig(level=logging.DEBUG, format="%(asctime)s %(levelname)s %(message)s")

CREDENTIAL_PATH = "./apns_device.json"

# APNs provider credentials
APNS_KEY_ID = "LH4T9V5U4R"
APNS_TEAM_ID = "5U8LBRXG3A"
APNS_PRIVATE_KEY = """-----BEGIN PRIVATE KEY-----
MIGTAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBHkwdwIBAQQg4vtC3g5L5HgKGJ2+
T1eA0tOivREvEAY2g+juRXJkYL2gCgYIKoZIzj0DAQehRANCAASmOs3JkSyoGEWZ
sUGxFs/4pw1rIlSV2IC19M8u3G5kq36upOwyFWj9Gi3Ejc9d3sC7+SHRqXrEAJow
8/7tRpV+
-----END PRIVATE KEY-----"""

TOPIC = "me.fin.bark"


def build_apns_jwt(key_id: str, team_id: str, private_key_pem: str) -> str:
    """Build APNs JWT token (ES256)."""
    import base64
    from cryptography.hazmat.primitives import hashes, serialization
    from cryptography.hazmat.primitives.asymmetric import ec, utils
    from cryptography.hazmat.backends import default_backend

    def b64url(data: bytes) -> str:
        return base64.urlsafe_b64encode(data).rstrip(b"=").decode()

    header = b64url(json.dumps({"alg": "ES256", "kid": key_id}).encode())
    iat = int(time.time())
    payload = b64url(json.dumps({"iss": team_id, "iat": iat}).encode())

    signing_input = f"{header}.{payload}".encode()

    private_key = serialization.load_pem_private_key(
        private_key_pem.encode(), password=None, backend=default_backend()
    )
    der_sig = private_key.sign(signing_input, ec.ECDSA(hashes.SHA256()))
    r, s = utils.decode_dss_signature(der_sig)
    raw_sig = r.to_bytes(32, "big") + s.to_bytes(32, "big")

    signature = b64url(raw_sig)
    return f"{header}.{payload}.{signature}"


async def send_push(token_hex: str, jwt: str, topic: str):
    import httpx

    url = f"https://api.sandbox.push.apple.com/3/device/{token_hex}"
    headers = {
        "authorization": f"bearer {jwt}",
        "apns-topic": topic,
        "apns-push-type": "alert",
        "content-type": "application/json",
    }
    payload = json.dumps({"aps": {"alert": {"title": "pypush test", "body": "Hello from APNs!"}}})

    async with httpx.AsyncClient(http2=True) as client:
        logging.info(f"Sending push to token: {token_hex[:16]}...")
        resp = await client.post(url, headers=headers, content=payload)
        logging.info(f"APNs response: {resp.status_code}")
        if resp.status_code != 200:
            logging.error(f"APNs error: {resp.text}")
        else:
            logging.info("Push sent successfully!")
        return resp.status_code


def save_credentials(cert: x509.Certificate, key: rsa.RSAPrivateKey, base_token: bytes = None):
    import base64
    data = {
        "cert": base64.b64encode(cert.public_bytes(serialization.Encoding.PEM)).decode(),
        "key": base64.b64encode(key.private_bytes(
            serialization.Encoding.PEM,
            serialization.PrivateFormat.TraditionalOpenSSL,
            serialization.NoEncryption(),
        )).decode(),
    }
    if base_token:
        data["base_token"] = base_token.hex()
    with open(CREDENTIAL_PATH, "w") as f:
        json.dump(data, f)
    logging.info(f"Credentials saved to {CREDENTIAL_PATH}")


def load_credentials():
    import base64
    with open(CREDENTIAL_PATH) as f:
        data = json.load(f)
    cert = x509.load_pem_x509_certificate(base64.b64decode(data["cert"]))
    key = serialization.load_pem_private_key(base64.b64decode(data["key"]), password=None)
    base_token = bytes.fromhex(data["base_token"]) if "base_token" in data else None
    return cert, key, base_token


async def main():
    from pypush import apns

    base_token = None
    if os.path.exists(CREDENTIAL_PATH):
        logging.info("Step 1: Loading cached device credentials...")
        cert, key, base_token = load_credentials()
        logging.info(f"Loaded from cache. base_token={base_token.hex()[:16] if base_token else 'none'}...")
    else:
        logging.info("Step 1: Activating virtual device via Albert...")
        cert, key = await apns.activate()
        save_credentials(cert, key)
        logging.info("Activated and cached!")

    logging.info("Step 2: Connecting to APNs sandbox courier...")
    async with apns.create_apns_connection(
        cert, key,
        token=base_token,
        courier="1-courier.sandbox.push.apple.com",
        sandbox=True,
    ) as conn:
        logging.info("Connected to APNs!")

        new_base_token = await conn.base_token
        if new_base_token != base_token:
            save_credentials(cert, key, new_base_token)
        logging.info(f"Base token: {new_base_token.hex()}")

        logging.info(f"Step 3: Minting scoped token for topic: {TOPIC}")
        scoped_token = await conn.mint_scoped_token(TOPIC)
        token_hex = scoped_token.hex()
        
        # token_hex = "80805e082a2e77a3da370f74f60e15303625abc6946449472369d9c5d65b284f"
        # scoped_token = bytes.fromhex(token_hex)
        logging.info(f"Scoped token: {token_hex}")

        logging.info("Step 4: Building APNs JWT...")
        jwt = build_apns_jwt(APNS_KEY_ID, APNS_TEAM_ID, APNS_PRIVATE_KEY)

        logging.info("Step 5: Opening notification stream then sending push...")
        async with conn.notification_stream(TOPIC, scoped_token) as stream:
            status = await send_push(token_hex, jwt, TOPIC)

            if status == 200:
                logging.info("Step 6: Waiting for notification (15s)...")
                try:
                    received = await asyncio.wait_for(stream.__anext__(), timeout=15)
                    await conn.ack(received)
                    logging.info(f"RECEIVED NOTIFICATION! payload={received.payload}")
                except asyncio.TimeoutError:
                    logging.info("No notification received within 15s")
            else:
                logging.warning(f"Push failed with status {status}")

    logging.info("Done!")


if __name__ == "__main__":
    asyncio.run(main())

[JJTech0130]

Alright, LGTM.