Skip to main content

wechat_protocol/
mmtls.rs

1//! Fresh pinned MMTLS handshake and authenticated shortlink exchange.
2use crate::crypto::*;
3use anyhow::{ensure, Result};
4use openssl::{memcmp, nid::Nid};
5use serde::{Deserialize, Serialize};
6use std::{
7    io::{Read, Write},
8    net::TcpStream,
9    time::{Duration, Instant},
10};
11use zeroize::Zeroize;
12
13pub const VERSION: u16 = 0xf104;
14pub struct Cursor<'a> {
15    bytes: &'a [u8],
16    offset: usize,
17}
18impl<'a> Cursor<'a> {
19    pub fn new(bytes: &'a [u8]) -> Self {
20        Self { bytes, offset: 0 }
21    }
22    pub fn take(&mut self, n: usize) -> Result<&'a [u8]> {
23        ensure!(n <= self.bytes.len() - self.offset, "Truncated structure");
24        let b = &self.bytes[self.offset..self.offset + n];
25        self.offset += n;
26        Ok(b)
27    }
28    pub fn uint(&mut self, n: usize) -> Result<u64> {
29        ensure!(n <= 8, "Integer too wide");
30        Ok(self.take(n)?.iter().fold(0, |a, &b| (a << 8) | b as u64))
31    }
32    pub fn sized(&mut self, width: usize) -> Result<&'a [u8]> {
33        let n = usize::try_from(self.uint(width)?)?;
34        self.take(n)
35    }
36    pub fn end(&self) -> Result<()> {
37        ensure!(self.offset == self.bytes.len(), "Trailing structure bytes");
38        Ok(())
39    }
40}
41pub fn section(b: &[u8]) -> Vec<u8> {
42    [(b.len() as u32).to_be_bytes().as_slice(), b].concat()
43}
44fn sized(b: &[u8]) -> Result<Vec<u8>> {
45    ensure!(b.len() <= u16::MAX as usize, "Section too large");
46    Ok([(b.len() as u16).to_be_bytes().as_slice(), b].concat())
47}
48#[derive(Clone)]
49pub struct Record {
50    pub header: [u8; 5],
51    pub body: Vec<u8>,
52}
53pub fn frame(kind: u8, body: &[u8]) -> Result<Vec<u8>> {
54    ensure!(
55        (0x15..=0x19).contains(&kind) && body.len() <= 65535,
56        "Invalid MMTLS record"
57    );
58    Ok([
59        vec![kind],
60        VERSION.to_be_bytes().to_vec(),
61        (body.len() as u16).to_be_bytes().to_vec(),
62        body.to_vec(),
63    ]
64    .concat())
65}
66pub fn nonce(iv: &[u8], sequence: u64) -> Result<[u8; 12]> {
67    ensure!(iv.len() == 12, "Invalid IV");
68    let mut b: [u8; 12] = iv.try_into()?;
69    for (i, v) in sequence.to_be_bytes().iter().enumerate() {
70        b[i + 4] ^= v;
71    }
72    Ok(b)
73}
74pub fn seal(kind: u8, plain: &[u8], key: &[u8], iv: &[u8], sequence: u64) -> Result<Vec<u8>> {
75    ensure!(plain.len() <= 65519, "MMTLS plaintext too large");
76    let header = frame(kind, &vec![0; plain.len() + 16])?[..5].to_vec();
77    let aad = [sequence.to_be_bytes().as_slice(), &header].concat();
78    let (cipher, tag) = gcm_seal(key, &nonce(iv, sequence)?, &aad, plain)?;
79    Ok([header, cipher, tag.to_vec()].concat())
80}
81pub fn open(record: &Record, key: &[u8], iv: &[u8], sequence: u64) -> Result<Vec<u8>> {
82    ensure!(record.body.len() >= 16, "Missing authentication tag");
83    let n = record.body.len() - 16;
84    gcm_open(
85        key,
86        &nonce(iv, sequence)?,
87        &[sequence.to_be_bytes().as_slice(), &record.header].concat(),
88        &record.body[..n],
89        &record.body[n..],
90    )
91}
92struct Traffic {
93    client_key: Vec<u8>,
94    server_key: Vec<u8>,
95    client_iv: Vec<u8>,
96    server_iv: Vec<u8>,
97}
98impl Drop for Traffic {
99    fn drop(&mut self) {
100        self.client_key.zeroize();
101        self.server_key.zeroize();
102        self.client_iv.zeroize();
103        self.server_iv.zeroize();
104    }
105}
106fn traffic(secret: &[u8], label: &str, transcript: &[u8]) -> Result<Traffic> {
107    let mut b = expand(
108        secret,
109        &[label.as_bytes(), &sha256(transcript)].concat(),
110        56,
111    )?;
112    let t = Traffic {
113        client_key: b[..16].to_vec(),
114        server_key: b[16..32].to_vec(),
115        client_iv: b[32..44].to_vec(),
116        server_iv: b[44..].to_vec(),
117    };
118    b.zeroize();
119    Ok(t)
120}
121fn signed_payload(bytes: &[u8], kind: u8) -> Result<Vec<u8>> {
122    let mut top = Cursor::new(bytes);
123    let mut body = Cursor::new(top.sized(4)?);
124    top.end()?;
125    ensure!(body.uint(1)? == kind as u64, "Unexpected handshake message");
126    let value = body.sized(2)?.to_vec();
127    body.end()?;
128    Ok(value)
129}
130
131#[derive(Clone, Serialize, Deserialize)]
132pub struct Ticket {
133    pub kind: u8,
134    pub lifetime: u32,
135    pub reserved: u32,
136    pub nonce: Vec<u8>,
137    pub opaque: Vec<u8>,
138}
139#[derive(Clone, Serialize, Deserialize)]
140pub struct ShortSession {
141    pub access_psk: Vec<u8>,
142    pub ticket: Ticket,
143}
144impl Drop for ShortSession {
145    fn drop(&mut self) {
146        self.access_psk.zeroize();
147        self.ticket.opaque.zeroize();
148        self.ticket.nonce.zeroize();
149    }
150}
151pub struct Session {
152    socket: TcpStream,
153    deadline: Instant,
154    received: usize,
155    receive_budget: usize,
156    pub short: ShortSession,
157    pub certificate_region: u32,
158    pub certificate_version: u32,
159    traffic: Traffic,
160    outgoing: u64,
161    incoming: u64,
162    record_pending: Vec<u8>,
163    application_pending: LonglinkDecoder,
164}
165
166fn host_valid(name: &str) -> bool {
167    !name.is_empty()
168        && name.len() <= 253
169        && name.is_ascii()
170        && name.split('.').all(|s| {
171            !s.is_empty()
172                && s.len() <= 63
173                && s.bytes()
174                    .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == b'-')
175                && !s.starts_with('-')
176                && !s.ends_with('-')
177        })
178        && (name.ends_with(".wechat.com") || name.ends_with(".weixin.qq.com"))
179}
180pub fn validate_host(name: &str) -> Result<()> {
181    ensure!(host_valid(name), "Unsupported native server host");
182    Ok(())
183}
184fn read_exact_deadline(
185    socket: &mut TcpStream,
186    mut bytes: &mut [u8],
187    deadline: Instant,
188) -> Result<()> {
189    while !bytes.is_empty() {
190        let remaining = deadline.saturating_duration_since(Instant::now());
191        ensure!(!remaining.is_zero(), "MMTLS deadline exceeded");
192        socket.set_read_timeout(Some(remaining))?;
193        match socket.read(bytes) {
194            Ok(0) => anyhow::bail!("MMTLS connection closed"),
195            Ok(n) => bytes = &mut bytes[n..],
196            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
197            Err(e) => return Err(e.into()),
198        }
199    }
200    Ok(())
201}
202fn read_record(
203    socket: &mut TcpStream,
204    deadline: Instant,
205    received: &mut usize,
206    budget: usize,
207) -> Result<Record> {
208    let mut h = [0; 5];
209    read_exact_deadline(socket, &mut h, deadline)?;
210    ensure!(
211        (0x15..=0x19).contains(&h[0]) && u16::from_be_bytes([h[1], h[2]]) == VERSION,
212        "Unexpected MMTLS header"
213    );
214    let n = u16::from_be_bytes([h[3], h[4]]) as usize;
215    *received = received
216        .checked_add(n + 5)
217        .ok_or_else(|| anyhow::anyhow!("Receive overflow"))?;
218    ensure!(*received <= budget, "MMTLS receive budget exceeded");
219    let mut body = vec![0; n];
220    read_exact_deadline(socket, &mut body, deadline)?;
221    ensure!(h[0] != 0x15, "Server returned MMTLS alert");
222    Ok(Record { header: h, body })
223}
224
225impl Session {
226    pub fn connect(host: &str, timeout: Duration, receive_budget: usize) -> Result<Self> {
227        validate_host(host)?;
228        ensure!(
229            (1024..=crate::MAX_MESSAGE_BYTES).contains(&receive_budget),
230            "Invalid receive budget"
231        );
232        ensure!(
233            !timeout.is_zero() && timeout <= Duration::from_secs(120),
234            "Invalid MMTLS timeout"
235        );
236        let deadline = Instant::now() + timeout;
237        // Resolver and connect are constrained by the caller's blocking worker. Socket
238        // read/write deadlines bound every subsequent protocol operation.
239        use std::net::ToSocketAddrs;
240        let addresses = (host, 80).to_socket_addrs()?;
241        let mut socket = None;
242        for addr in addresses {
243            let remaining = deadline.saturating_duration_since(Instant::now());
244            if remaining.is_zero() {
245                break;
246            }
247            if let Ok(s) = TcpStream::connect_timeout(&addr, remaining) {
248                socket = Some(s);
249                break;
250            }
251        }
252        let mut socket = socket.ok_or_else(|| anyhow::anyhow!("Native connection failed"))?;
253        socket.set_write_timeout(Some(timeout))?;
254        let keys = [
255            Ecdh::generate(Nid::X9_62_PRIME256V1)?,
256            Ecdh::generate(Nid::X9_62_PRIME256V1)?,
257        ];
258        let mut shares = vec![];
259        for (i, k) in keys.iter().enumerate() {
260            let point = k.public()?;
261            shares.extend(section(
262                &[((i + 5) as u32).to_be_bytes().to_vec(), sized(&point)?].concat(),
263            ));
264        }
265        let mut extra = vec![0; 4];
266        extra.push(2);
267        extra.extend(5u32.to_be_bytes());
268        extra.extend(1000u32.to_be_bytes());
269        let extension = section(&[0x10u16.to_be_bytes().to_vec(), vec![2], shares, extra].concat());
270        let hello = section(
271            &[
272                vec![1],
273                VERSION.to_le_bytes().to_vec(),
274                vec![1],
275                0xc02bu16.to_be_bytes().to_vec(),
276                random::<32>()?.to_vec(),
277                crate::wire::unix_time()?.to_be_bytes().to_vec(),
278                section(&[vec![1], extension].concat()),
279            ]
280            .concat(),
281        );
282        socket.write_all(&frame(0x16, &hello)?)?;
283        let mut received = 0;
284        let server_record = read_record(&mut socket, deadline, &mut received, receive_budget)?;
285        ensure!(server_record.header[0] == 0x16, "Expected handshake");
286        let server = parse_server_hello(&server_record.body)?;
287        let index = server
288            .key_id
289            .checked_sub(5)
290            .ok_or_else(|| anyhow::anyhow!("Unoffered key"))? as usize;
291        ensure!(index < 2, "Unoffered key");
292        let mut x = keys[index].derive(&server.point)?;
293        let first = x.iter().position(|v| *v != 0).unwrap_or(x.len() - 1);
294        let mut secret = sha256(&x[first..]);
295        x.zeroize();
296        let mut transcript = [hello, server_record.body].concat();
297        let handshake = traffic(&secret, "handshake key expansion", &transcript)?;
298        let r = read_record(&mut socket, deadline, &mut received, receive_budget)?;
299        ensure!(r.header[0] == 0x16, "Expected signature record");
300        let signature_body = open(&r, &handshake.server_key, &handshake.server_iv, 1)?;
301        let signature = signed_payload(&signature_body, 0x0f)?;
302        let pins: serde_json::Value =
303            serde_json::from_str(include_str!("../fixtures/mmtls-public-keys.json"))?;
304        let pin = pins["profiles"]
305            .as_array()
306            .unwrap()
307            .iter()
308            .find(|p| p["version"].as_u64() == Some(server.version as u64))
309            .ok_or_else(|| anyhow::anyhow!("Unsupported MMTLS certificate version"))?;
310        ensure!([5, 1000].contains(&server.version), "Unoffered MMTLS pin");
311        verify_signature(
312            pin["ecdsaPEM"].as_str().unwrap(),
313            &sha256(&transcript),
314            &signature,
315        )?;
316        transcript.extend(signature_body);
317        let r = read_record(&mut socket, deadline, &mut received, receive_budget)?;
318        ensure!(r.header[0] == 0x16, "Expected ticket record");
319        let ticket_body = open(&r, &handshake.server_key, &handshake.server_iv, 2)?;
320        let tickets = parse_tickets(&ticket_body)?;
321        let access_psk = expand(
322            &secret,
323            &[b"PSK_ACCESS".as_slice(), &sha256(&transcript)].concat(),
324            32,
325        )?;
326        transcript.extend(ticket_body);
327        let r = read_record(&mut socket, deadline, &mut received, receive_budget)?;
328        ensure!(r.header[0] == 0x16, "Expected Finished");
329        let finish = signed_payload(
330            &open(&r, &handshake.server_key, &handshake.server_iv, 3)?,
331            0x14,
332        )?;
333        let expected = hmac(
334            &expand(&secret, b"server finished", 32)?,
335            &sha256(&transcript),
336        )?;
337        ensure!(
338            finish.len() == expected.len() && memcmp::eq(&finish, &expected),
339            "Server Finished verification failed"
340        );
341        let client_finish = hmac(
342            &expand(&secret, b"client finished", 32)?,
343            &sha256(&transcript),
344        )?;
345        socket.write_all(&seal(
346            0x16,
347            &section(&[vec![0x14], sized(&client_finish)?].concat()),
348            &handshake.client_key,
349            &handshake.client_iv,
350            1,
351        )?)?;
352        let mut expanded = expand(
353            &secret,
354            &[b"expanded secret".as_slice(), &sha256(&transcript)].concat(),
355            32,
356        )?;
357        let application = traffic(&expanded, "application data key expansion", &transcript)?;
358        secret.zeroize();
359        expanded.zeroize();
360        let noop = pack_longlink(6, u32::MAX, &[])?;
361        socket.write_all(&seal(
362            0x17,
363            &noop,
364            &application.client_key,
365            &application.client_iv,
366            2,
367        )?)?;
368        let r = read_record(&mut socket, deadline, &mut received, receive_budget)?;
369        ensure!(r.header[0] == 0x17, "Expected NOOP reply");
370        let p = unpack_longlink(&open(
371            &r,
372            &application.server_key,
373            &application.server_iv,
374            4,
375        )?)?;
376        ensure!(
377            p.command == 0x3b9aca06 && p.sequence == u32::MAX && p.payload.is_empty(),
378            "Unexpected NOOP reply"
379        );
380        Ok(Self {
381            socket,
382            deadline,
383            received,
384            receive_budget,
385            short: ShortSession {
386                access_psk,
387                ticket: tickets[0].clone(),
388            },
389            certificate_region: server.region,
390            certificate_version: server.version,
391            traffic: application,
392            outgoing: 3,
393            incoming: 5,
394            record_pending: Vec::new(),
395            application_pending: LonglinkDecoder {
396                pending: Vec::new(),
397                max_payload: receive_budget,
398            },
399        })
400    }
401    pub fn send_application(&mut self, command: u32, sequence: u32, payload: &[u8]) -> Result<()> {
402        let bytes = seal(
403            0x17,
404            &pack_longlink(command, sequence, payload)?,
405            &self.traffic.client_key,
406            &self.traffic.client_iv,
407            self.outgoing,
408        )?;
409        self.socket.write_all(&bytes)?;
410        self.outgoing = self
411            .outgoing
412            .checked_add(1)
413            .ok_or_else(|| anyhow::anyhow!("Sequence exhausted"))?;
414        Ok(())
415    }
416    pub fn receive_application(&mut self) -> Result<Longlink> {
417        let r = read_record(
418            &mut self.socket,
419            self.deadline,
420            &mut self.received,
421            self.receive_budget,
422        )?;
423        ensure!(r.header[0] == 0x17, "Expected application record");
424        let plain = open(
425            &r,
426            &self.traffic.server_key,
427            &self.traffic.server_iv,
428            self.incoming,
429        )?;
430        self.incoming = self
431            .incoming
432            .checked_add(1)
433            .ok_or_else(|| anyhow::anyhow!("Sequence exhausted"))?;
434        unpack_longlink(&plain)
435    }
436
437    /// A bounded idle wait on an established connection. Partial encrypted
438    /// records and application packets survive a timeout; MMTLS record
439    /// boundaries need not coincide with longlink packet boundaries.
440    pub fn poll_application(&mut self, timeout: Duration) -> Result<Option<Longlink>> {
441        ensure!(
442            !timeout.is_zero() && timeout <= Duration::from_secs(120),
443            "Invalid longlink poll timeout"
444        );
445        let deadline = Instant::now() + timeout;
446        loop {
447            if let Some(packet) = self.application_pending.next_packet()? {
448                return Ok(Some(packet));
449            }
450            if self.record_pending.len() >= 5 {
451                let h: [u8; 5] = self.record_pending[..5].try_into()?;
452                ensure!(
453                    (0x15..=0x19).contains(&h[0]) && u16::from_be_bytes([h[1], h[2]]) == VERSION,
454                    "Unexpected MMTLS header"
455                );
456                let n = u16::from_be_bytes([h[3], h[4]]) as usize;
457                if self.record_pending.len() >= n + 5 {
458                    ensure!(h[0] == 0x17, "Unexpected persistent MMTLS record");
459                    let record = Record {
460                        header: h,
461                        body: self.record_pending[5..n + 5].to_vec(),
462                    };
463                    let plain = open(
464                        &record,
465                        &self.traffic.server_key,
466                        &self.traffic.server_iv,
467                        self.incoming,
468                    )?;
469                    self.incoming = self
470                        .incoming
471                        .checked_add(1)
472                        .ok_or_else(|| anyhow::anyhow!("Sequence exhausted"))?;
473                    self.record_pending.drain(..n + 5);
474                    self.application_pending.feed(&plain)?;
475                    continue;
476                }
477            }
478            let remaining = deadline.saturating_duration_since(Instant::now());
479            if remaining.is_zero() {
480                return Ok(None);
481            }
482            self.socket.set_read_timeout(Some(remaining))?;
483            let mut bytes = [0u8; 65540];
484            match self.socket.read(&mut bytes) {
485                Ok(0) => anyhow::bail!("MMTLS connection closed"),
486                Ok(n) => {
487                    // At most one incomplete record precedes a socket read.
488                    ensure!(
489                        self.record_pending.len() + n <= 131080,
490                        "MMTLS pending record limit"
491                    );
492                    self.record_pending.extend_from_slice(&bytes[..n]);
493                }
494                Err(e)
495                    if matches!(
496                        e.kind(),
497                        std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
498                    ) =>
499                {
500                    return Ok(None)
501                }
502                Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
503                Err(e) => return Err(e.into()),
504            }
505        }
506    }
507}
508struct ServerHello {
509    key_id: u32,
510    point: Vec<u8>,
511    region: u32,
512    version: u32,
513}
514fn parse_server_hello(bytes: &[u8]) -> Result<ServerHello> {
515    let mut top = Cursor::new(bytes);
516    let mut b = Cursor::new(top.sized(4)?);
517    top.end()?;
518    ensure!(
519        b.uint(1)? == 2 && b.take(2)? == VERSION.to_le_bytes() && b.uint(2)? == 0xc02b,
520        "Unsupported ServerHello"
521    );
522    b.take(32)?;
523    let mut e = Cursor::new(b.sized(4)?);
524    b.end()?;
525    let count = e.uint(1)?;
526    let mut seen = std::collections::HashSet::new();
527    let mut share = None;
528    let mut cert = None;
529    for _ in 0..count {
530        let mut x = Cursor::new(e.sized(4)?);
531        let kind = x.uint(2)?;
532        ensure!(seen.insert(kind), "Duplicate extension");
533        match kind {
534            0x11 => {
535                let id = x.uint(4)? as u32;
536                let point = x.sized(2)?.to_vec();
537                ensure!(point.len() == 65 && point[0] == 4, "Invalid ECDH point");
538                x.end()?;
539                share = Some((id, point));
540            }
541            0x13 => {
542                let region = x.uint(4)? as u32;
543                let version = x.uint(4)? as u32;
544                x.end()?;
545                cert = Some((region, version));
546            }
547            _ => {}
548        }
549    }
550    e.end()?;
551    let (key_id, point) = share.ok_or_else(|| anyhow::anyhow!("Missing ECDH selection"))?;
552    let (region, version) = cert.ok_or_else(|| anyhow::anyhow!("Missing certificate selection"))?;
553    Ok(ServerHello {
554        key_id,
555        point,
556        region,
557        version,
558    })
559}
560fn parse_tickets(input: &[u8]) -> Result<Vec<Ticket>> {
561    let mut top = Cursor::new(input);
562    let mut b = Cursor::new(top.sized(4)?);
563    top.end()?;
564    ensure!(b.uint(1)? == 4, "Expected tickets");
565    let count = b.uint(1)?;
566    ensure!((1..=4).contains(&count), "Invalid ticket count");
567    let mut out = vec![];
568    for _ in 0..count {
569        let mut t = Cursor::new(b.sized(4)?);
570        let kind = t.uint(1)? as u8;
571        let lifetime = t.uint(4)? as u32;
572        t.sized(2)?;
573        let reserved = t.uint(4)? as u32;
574        let nonce = t.sized(2)?.to_vec();
575        let opaque = t.sized(2)?.to_vec();
576        t.end()?;
577        ensure!(
578            nonce.len() == 12 && !opaque.is_empty(),
579            "Invalid shortlink ticket"
580        );
581        out.push(Ticket {
582            kind,
583            lifetime,
584            reserved,
585            nonce,
586            opaque,
587        });
588    }
589    b.end()?;
590    Ok(out)
591}
592pub struct Longlink {
593    pub command: u32,
594    pub sequence: u32,
595    pub payload: Vec<u8>,
596}
597struct LonglinkDecoder {
598    pending: Vec<u8>,
599    max_payload: usize,
600}
601impl Default for LonglinkDecoder {
602    fn default() -> Self {
603        Self {
604            pending: Vec::new(),
605            max_payload: crate::MAX_MESSAGE_BYTES,
606        }
607    }
608}
609impl LonglinkDecoder {
610    fn feed(&mut self, bytes: &[u8]) -> Result<()> {
611        ensure!(
612            self.pending.len().saturating_add(bytes.len()) <= self.max_payload + 65540,
613            "Longlink pending packet limit"
614        );
615        self.pending.extend_from_slice(bytes);
616        Ok(())
617    }
618    fn next_packet(&mut self) -> Result<Option<Longlink>> {
619        if self.pending.len() < 4 {
620            return Ok(None);
621        }
622        let n = u32::from_be_bytes(self.pending[..4].try_into()?) as usize;
623        ensure!(
624            (16..=self.max_payload + 16).contains(&n),
625            "Invalid longlink packet length"
626        );
627        if self.pending.len() >= 8 {
628            ensure!(
629                self.pending[4..8] == [0, 16, 0, 1],
630                "Invalid longlink envelope"
631            );
632        }
633        if self.pending.len() < n {
634            return Ok(None);
635        }
636        let packet = unpack_longlink(&self.pending[..n])?;
637        self.pending.drain(..n);
638        Ok(Some(packet))
639    }
640}
641pub fn pack_longlink(command: u32, sequence: u32, payload: &[u8]) -> Result<Vec<u8>> {
642    ensure!(payload.len() <= 65503, "Longlink body too large");
643    Ok([
644        ((16 + payload.len()) as u32).to_be_bytes().to_vec(),
645        0x10u16.to_be_bytes().to_vec(),
646        1u16.to_be_bytes().to_vec(),
647        command.to_be_bytes().to_vec(),
648        sequence.to_be_bytes().to_vec(),
649        payload.to_vec(),
650    ]
651    .concat())
652}
653pub fn unpack_longlink(input: &[u8]) -> Result<Longlink> {
654    let mut c = Cursor::new(input);
655    ensure!(
656        c.uint(4)? == input.len() as u64 && c.uint(2)? == 0x10 && c.uint(2)? == 1,
657        "Invalid longlink envelope"
658    );
659    let command = c.uint(4)? as u32;
660    let sequence = c.uint(4)? as u32;
661    let payload = c.take(input.len() - 16)?.to_vec();
662    Ok(Longlink {
663        command,
664        sequence,
665        payload,
666    })
667}
668
669pub fn short_exchange(
670    session: &ShortSession,
671    host: &str,
672    path: &str,
673    payload: &[u8],
674    timeout: Duration,
675) -> Result<Vec<u8>> {
676    validate_host(host)?;
677    ensure!(
678        path.starts_with("/cgi-bin/micromsg-bin/")
679            && path
680                .bytes()
681                .all(|b| b.is_ascii_alphanumeric() || b == b'/' || b == b'-'),
682        "Unsupported native route"
683    );
684    ensure!(
685        payload.len() <= crate::MAX_MESSAGE_BYTES,
686        "Payload too large"
687    );
688    let (hello, extensions) = early_hello(&session.ticket)?;
689    let keys = expand(
690        &session.access_psk,
691        &[b"early data key expansion".as_slice(), &sha256(&hello)].concat(),
692        28,
693    )?;
694    let route = section(
695        &[
696            sized(path.as_bytes())?,
697            sized(host.as_bytes())?,
698            section(payload),
699        ]
700        .concat(),
701    );
702    let wire = [
703        frame(0x19, &hello)?,
704        seal(0x19, &extensions, &keys[..16], &keys[16..], 1)?,
705        seal(0x17, &route, &keys[..16], &keys[16..], 2)?,
706        seal(0x15, &section(&[0, 1, 1]), &keys[..16], &keys[16..], 3)?,
707    ]
708    .concat();
709    let client = reqwest::blocking::Client::builder()
710        .timeout(timeout)
711        .connect_timeout(timeout.min(Duration::from_secs(10)))
712        .redirect(reqwest::redirect::Policy::none())
713        .no_proxy()
714        .build()?;
715    let mut response = client
716        .post(format!(
717            "http://{host}/mmtls/{}",
718            hex::encode(random::<4>()?)
719        ))
720        .header("Content-Type", "application/octet-stream")
721        .header("User-Agent", "MicroMessenger Client")
722        .header("Upgrade", "mmtls")
723        .header("Accept", "*/*")
724        .header("Cache-Control", "no-cache")
725        .header("Connection", "close")
726        .body(wire)
727        .send()?;
728    ensure!(
729        response.status() == 200,
730        "Native shortlink HTTP status {}",
731        response.status().as_u16()
732    );
733    let mut bytes = Vec::new();
734    response
735        .by_ref()
736        .take((crate::MAX_MESSAGE_BYTES + 1) as u64)
737        .read_to_end(&mut bytes)?;
738    ensure!(
739        bytes.len() <= crate::MAX_MESSAGE_BYTES,
740        "Shortlink receive budget exceeded"
741    );
742    decode_short_reply(&bytes, session, &hello, &extensions)
743}
744fn early_hello(ticket: &Ticket) -> Result<(Vec<u8>, Vec<u8>)> {
745    ensure!(
746        ticket.nonce.len() == 12 && !ticket.opaque.is_empty(),
747        "Invalid short ticket"
748    );
749    let now = crate::wire::unix_time()?;
750    let encoded = [
751        vec![ticket.kind],
752        ticket.lifetime.to_be_bytes().to_vec(),
753        sized(&[])?,
754        ticket.reserved.to_be_bytes().to_vec(),
755        sized(&ticket.nonce)?,
756        sized(&ticket.opaque)?,
757    ]
758    .concat();
759    let ext = section(&[0x0fu16.to_be_bytes().to_vec(), vec![1], section(&encoded)].concat());
760    let hello = section(
761        &[
762            vec![1],
763            VERSION.to_le_bytes().to_vec(),
764            vec![1],
765            0xa8u16.to_be_bytes().to_vec(),
766            random::<32>()?.to_vec(),
767            now.to_be_bytes().to_vec(),
768            section(&[vec![1], ext].concat()),
769        ]
770        .concat(),
771    );
772    let extensions = section(
773        &[
774            vec![8],
775            section(
776                &[
777                    vec![1],
778                    section(&[0x12u16.to_be_bytes().to_vec(), now.to_be_bytes().to_vec()].concat()),
779                ]
780                .concat(),
781            ),
782        ]
783        .concat(),
784    );
785    Ok((hello, extensions))
786}
787fn decode_short_reply(
788    input: &[u8],
789    session: &ShortSession,
790    hello: &[u8],
791    extensions: &[u8],
792) -> Result<Vec<u8>> {
793    let mut c = Cursor::new(input);
794    let mut records = vec![];
795    while c.offset < input.len() {
796        let h: [u8; 5] = c.take(5)?.try_into()?;
797        ensure!(
798            (0x15..=0x19).contains(&h[0]) && u16::from_be_bytes([h[1], h[2]]) == VERSION,
799            "Invalid shortlink header"
800        );
801        records.push(Record {
802            header: h,
803            body: c.take(u16::from_be_bytes([h[3], h[4]]) as usize)?.to_vec(),
804        });
805    }
806    ensure!(
807        records.len() >= 4
808            && records[0].header[0] == 0x16
809            && records[1].header[0] == 0x16
810            && records.last().unwrap().header[0] == 0x15
811            && records[2..records.len() - 1]
812                .iter()
813                .all(|r| r.header[0] == 0x17),
814        "Invalid shortlink sequence"
815    );
816    let server = &records[0];
817    let mut top = Cursor::new(&server.body);
818    let mut b = Cursor::new(top.sized(4)?);
819    top.end()?;
820    ensure!(
821        b.uint(1)? == 2 && b.take(2)? == VERSION.to_le_bytes() && b.uint(2)? == 0xa8,
822        "Unsupported shortlink selection"
823    );
824    b.take(32)?;
825    let mut e = Cursor::new(b.sized(4)?);
826    b.end()?;
827    ensure!(e.uint(1)? == 0, "Unsupported shortlink extensions");
828    e.end()?;
829    let transcript = [hello, extensions, &server.body].concat();
830    let keys = expand(
831        &session.access_psk,
832        &[b"handshake key expansion".as_slice(), &sha256(&transcript)].concat(),
833        28,
834    )?;
835    let finish = signed_payload(&open(&records[1], &keys[..16], &keys[16..], 1)?, 0x14)?;
836    let expected = hmac(
837        &expand(&session.access_psk, b"server finished", 32)?,
838        &sha256(&transcript),
839    )?;
840    ensure!(
841        finish.len() == expected.len() && memcmp::eq(&finish, &expected),
842        "Shortlink Finished failed"
843    );
844    let mut plain = vec![];
845    for (i, r) in records[2..records.len() - 1].iter().enumerate() {
846        plain.extend(open(r, &keys[..16], &keys[16..], i as u64 + 2)?);
847        ensure!(
848            plain.len() <= crate::MAX_MESSAGE_BYTES,
849            "Shortlink plaintext too large"
850        );
851    }
852    ensure!(
853        open(
854            records.last().unwrap(),
855            &keys[..16],
856            &keys[16..],
857            records.len() as u64 - 1
858        )? == section(&[0, 1, 1]),
859        "Unexpected shortlink close"
860    );
861    Ok(plain)
862}
863
864#[cfg(test)]
865mod tests {
866    use super::*;
867    #[test]
868    fn longlink_frames_can_split_or_share_records_and_reject_bad_lengths_early() {
869        let a = pack_longlink(24, 0, &6u32.to_be_bytes()).unwrap();
870        let b = pack_longlink(1000000238, 23, &[2; 256]).unwrap();
871        let both = [a.as_slice(), b.as_slice()].concat();
872        for cut in 0..a.len() {
873            let mut decoder = LonglinkDecoder::default();
874            decoder.feed(&both[..cut]).unwrap();
875            assert!(decoder.next_packet().unwrap().is_none());
876            decoder.feed(&both[cut..]).unwrap();
877            assert_eq!(decoder.next_packet().unwrap().unwrap().command, 24);
878            assert_eq!(decoder.next_packet().unwrap().unwrap().sequence, 23);
879            assert!(decoder.next_packet().unwrap().is_none());
880        }
881        for n in [0u32, 15, (crate::MAX_MESSAGE_BYTES + 17) as u32, u32::MAX] {
882            let mut decoder = LonglinkDecoder::default();
883            decoder.feed(&n.to_be_bytes()).unwrap();
884            assert!(decoder.next_packet().is_err());
885        }
886        let mut invalid = a.clone();
887        invalid[7] = 2;
888        let mut decoder = LonglinkDecoder::default();
889        decoder.feed(&invalid[..8]).unwrap();
890        assert!(decoder.next_packet().is_err());
891        let mut limited = LonglinkDecoder {
892            pending: Vec::new(),
893            max_payload: 128,
894        };
895        limited.feed(&b[..4]).unwrap();
896        assert!(limited.next_packet().is_err());
897        let payload = vec![7; 200_000];
898        let large = [
899            (payload.len() as u32 + 16).to_be_bytes().to_vec(),
900            vec![0, 16, 0, 1],
901            1000000121u32.to_be_bytes().to_vec(),
902            9u32.to_be_bytes().to_vec(),
903            payload.clone(),
904        ]
905        .concat();
906        let mut decoder = LonglinkDecoder::default();
907        for part in large.chunks(20_000) {
908            decoder.feed(part).unwrap();
909        }
910        assert_eq!(decoder.next_packet().unwrap().unwrap().payload, payload);
911    }
912    #[test]
913    fn encrypted_idle_timeout_keeps_partial_record_and_packet_bytes() {
914        use std::{net::TcpListener, sync::mpsc, thread};
915        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
916        let socket = TcpStream::connect(listener.local_addr().unwrap()).unwrap();
917        let (mut peer, _) = listener.accept().unwrap();
918        let first = pack_longlink(24, 0, &6u32.to_be_bytes()).unwrap();
919        let second = pack_longlink(1000000006, u32::MAX, &[]).unwrap();
920        let key = vec![1; 16];
921        let iv = vec![2; 12];
922        let r1 = seal(0x17, &first[..10], &key, &iv, 5).unwrap();
923        let r2 = seal(0x17, &[first[10..].to_vec(), second].concat(), &key, &iv, 6).unwrap();
924        let mut corrupt = seal(0x17, &first, &key, &iv, 7).unwrap();
925        *corrupt.last_mut().unwrap() ^= 1;
926        let (sent, observed) = mpsc::channel();
927        let (resume, waiting) = mpsc::channel();
928        let writer = thread::spawn(move || {
929            peer.write_all(&r1[..3]).unwrap();
930            sent.send(()).unwrap();
931            waiting.recv().unwrap();
932            peer.write_all(&r1[3..]).unwrap();
933            peer.write_all(&r2).unwrap();
934            peer.write_all(&corrupt).unwrap();
935        });
936        let mut session = Session {
937            socket,
938            deadline: Instant::now(),
939            received: 0,
940            receive_budget: 1024,
941            short: ShortSession {
942                access_psk: vec![],
943                ticket: Ticket {
944                    kind: 1,
945                    lifetime: 0,
946                    reserved: 0,
947                    nonce: vec![],
948                    opaque: vec![],
949                },
950            },
951            certificate_region: 0,
952            certificate_version: 0,
953            traffic: Traffic {
954                client_key: key.clone(),
955                server_key: key,
956                client_iv: iv.clone(),
957                server_iv: iv,
958            },
959            outgoing: 3,
960            incoming: 5,
961            record_pending: Vec::new(),
962            application_pending: LonglinkDecoder::default(),
963        };
964        observed.recv().unwrap();
965        assert!(session
966            .poll_application(Duration::from_millis(10))
967            .unwrap()
968            .is_none());
969        assert_eq!(session.record_pending.len(), 3);
970        resume.send(()).unwrap();
971        assert_eq!(
972            session
973                .poll_application(Duration::from_secs(1))
974                .unwrap()
975                .unwrap()
976                .command,
977            24
978        );
979        assert_eq!(
980            session
981                .poll_application(Duration::from_secs(1))
982                .unwrap()
983                .unwrap()
984                .sequence,
985            u32::MAX
986        );
987        assert!(session.poll_application(Duration::from_secs(1)).is_err());
988        assert_eq!(session.incoming, 7);
989        writer.join().unwrap();
990    }
991    #[test]
992    fn published_mmtls_vector() {
993        let key = hex::decode("29b37affd98238fb0370c564d8fa992e").unwrap();
994        let wire_nonce = hex::decode("36ee037a258fbb83ca1f9762").unwrap();
995        let base = nonce(&wire_nonce, 1).unwrap();
996        let plain = hex::decode(
997            "0000002314002009d6fc43a65a48dd64e3b6846562d13752f5abc2cb254539906ecd429df71082",
998        )
999        .unwrap();
1000        let expected=hex::decode("16f1040037e4bdea652cd80a273fd2ef549df4c87c1083c23f658fcc2d4152d163bc809f35a56fd0f4a62caa138d765a086eec02e21704892a981cf7").unwrap();
1001        assert_eq!(seal(0x16, &plain, &key, &base, 1).unwrap(), expected);
1002        let r = Record {
1003            header: expected[..5].try_into().unwrap(),
1004            body: expected[5..].to_vec(),
1005        };
1006        assert_eq!(open(&r, &key, &base, 1).unwrap(), plain);
1007        assert!(open(&r, &key, &base, 2).is_err());
1008    }
1009    #[test]
1010    fn hosts_restricted() {
1011        for s in [
1012            "localhost",
1013            "127.0.0.1",
1014            "long.weixin.qq.com.attacker.org",
1015            "-x.wechat.com",
1016            "x.wechat.com:80",
1017        ] {
1018            assert!(validate_host(s).is_err());
1019        }
1020        assert!(validate_host("long.weixin.qq.com").is_ok());
1021    }
1022    #[test]
1023    fn malformed_envelopes() {
1024        assert!(unpack_longlink(&[]).is_err());
1025        assert!(parse_server_hello(&[0; 4]).is_err());
1026        let mut v = pack_longlink(6, 123, &[1, 2]).unwrap();
1027        v[3] ^= 1;
1028        assert!(unpack_longlink(&v).is_err());
1029    }
1030}