Skip to main content

autoit/
token.rs

1//! EA06 tokenized script parsing.
2//!
3//! Decodes the opcode-prefixed token stream of a compiled AutoIt (EA06) script
4//! into a [`TokenStream`] via [`parse`], and can re-render that stream back to
5//! source-like text with [`TokenStream::render_source`]. Keyword, function, and
6//! macro names are resolved and case-canonicalized through the lookup tables in
7//! the `tables` submodule.
8
9mod tables;
10
11use crate::{Error, util};
12
13const OPCODE_KEYWORD_ID: u8 = 0x00;
14const OPCODE_FUNCTION_ID: u8 = 0x01;
15const OPCODE_U32: u8 = 0x05;
16const OPCODE_U64: u8 = 0x10;
17const OPCODE_F64: u8 = 0x20;
18const OPCODE_KEYWORD_STRING: u8 = 0x30;
19const OPCODE_FUNCTION_STRING: u8 = 0x31;
20const OPCODE_MACRO_STRING: u8 = 0x32;
21const OPCODE_VARIABLE_STRING: u8 = 0x33;
22const OPCODE_BARE_STRING: u8 = 0x34;
23const OPCODE_PROPERTY_STRING: u8 = 0x35;
24const OPCODE_QUOTED_STRING: u8 = 0x36;
25const OPCODE_RAW_STRING: u8 = 0x37;
26const OPCODE_LINE_END: u8 = 0x7f;
27
28/// Parsed token stream.
29#[derive(Debug, Clone, PartialEq)]
30pub struct TokenStream {
31    line_count: u32,
32    tokens: Vec<Token>,
33}
34
35impl TokenStream {
36    /// Returns the line count declared in the stream header.
37    ///
38    /// # Returns
39    ///
40    /// The number of source lines [`parse`] read from the leading `u32` count,
41    /// which equals the number of [`Token::LineEnd`] tokens in the stream.
42    #[must_use]
43    pub const fn line_count(&self) -> u32 {
44        self.line_count
45    }
46
47    /// Returns the parsed tokens in stream order.
48    ///
49    /// # Returns
50    ///
51    /// A slice of every [`Token`] read from the input, including the
52    /// [`Token::LineEnd`] markers that terminate each source line.
53    #[must_use]
54    pub fn tokens(&self) -> &[Token] {
55        self.tokens.as_slice()
56    }
57
58    /// Renders source-like AutoIt text from the token stream.
59    ///
60    /// Accumulates tokens per line until a [`Token::LineEnd`] is reached, then
61    /// emits the line prefixed with tab indentation and terminated by `\r\n`.
62    /// Indentation is tracked across lines: `line_indent` computes the indent
63    /// to apply to the current line (dedenting block-closing keywords) and
64    /// `next_indent` computes the level for the following line (indenting
65    /// after block-opening keywords). Each non-`LineEnd` token contributes its
66    /// `display_text`, joined by `render_line`.
67    ///
68    /// # Returns
69    ///
70    /// The reconstructed script text with CRLF line endings and tab indentation.
71    #[must_use]
72    pub fn render_source(&self) -> String {
73        let mut out = String::new();
74        let mut line = Vec::new();
75        let mut indent = 0usize;
76        for token in &self.tokens {
77            match token {
78                Token::LineEnd => {
79                    let line_indent = line_indent(indent, line.as_slice());
80                    out.push_str("\t".repeat(line_indent).as_str());
81                    out.push_str(render_line(line.as_slice()).as_str());
82                    out.push_str("\r\n");
83                    indent = next_indent(indent, line.as_slice());
84                    line.clear();
85                }
86                other => line.push(other.display_text()),
87            }
88        }
89        out
90    }
91}
92
93/// Token parser error.
94#[derive(Debug, Clone, Copy, PartialEq, Eq)]
95pub enum TokenError {
96    /// Input ended before a token could be read.
97    Truncated,
98    /// Token string bytes were malformed.
99    BadString,
100}
101
102/// Parsed token.
103#[derive(Debug, Clone, PartialEq)]
104pub enum Token {
105    /// Keyword token.
106    Keyword(String),
107    /// Unknown keyword id.
108    UnknownKeywordId(i32),
109    /// Function token.
110    Function(String),
111    /// Unknown function id.
112    UnknownFunctionId(i32),
113    /// Macro token, without the leading `@`.
114    Macro(String),
115    /// Variable token, without the leading `$`.
116    Variable(String),
117    /// Bare identifier/string token.
118    BareString(String),
119    /// Property/member token, without the leading `.`.
120    Property(String),
121    /// Quoted string literal, unescaped.
122    QuotedString(String),
123    /// Raw string token.
124    RawString(String),
125    /// Unsigned 32-bit integer literal.
126    U32(u32),
127    /// Unsigned 64-bit integer literal.
128    U64(u64),
129    /// Floating-point literal.
130    F64(f64),
131    /// Operator or punctuation token.
132    Operator(&'static str),
133    /// Unknown opcode.
134    UnknownOpcode(u8),
135    /// End-of-line token.
136    LineEnd,
137}
138
139impl Token {
140    /// Renders this token to its source-like display string.
141    ///
142    /// Each variant maps to its concrete syntax: sigils are re-attached
143    /// (`@` for [`Token::Macro`], `$` for [`Token::Variable`], `.` for
144    /// [`Token::Property`]), [`Token::QuotedString`] is wrapped in double
145    /// quotes with embedded `"` doubled, and unresolved variants render as
146    /// angle-bracketed placeholders (e.g. `<keyword:5>`, `<opcode:0xff>`).
147    /// [`Token::LineEnd`] renders as the empty string.
148    ///
149    /// # Returns
150    ///
151    /// The display text for this token.
152    fn display_text(&self) -> String {
153        match self {
154            Self::Keyword(value)
155            | Self::Function(value)
156            | Self::BareString(value)
157            | Self::RawString(value) => value.clone(),
158            Self::UnknownKeywordId(value) => format!("<keyword:{value}>"),
159            Self::UnknownFunctionId(value) => format!("<function:{value}>"),
160            Self::Macro(value) => format!("@{value}"),
161            Self::Variable(value) => format!("${value}"),
162            Self::Property(value) => format!(".{value}"),
163            Self::QuotedString(value) => format!("\"{}\"", value.replace('"', "\"\"")),
164            Self::U32(value) => value.to_string(),
165            Self::U64(value) => value.to_string(),
166            Self::F64(value) => value.to_string(),
167            Self::Operator(value) => (*value).to_string(),
168            Self::UnknownOpcode(value) => format!("<opcode:{value:#x}>"),
169            Self::LineEnd => String::new(),
170        }
171    }
172}
173
174/// Parses a tokenized AutoIt (EA06) script into a [`TokenStream`].
175///
176/// Reads the leading little-endian `u32` line count, then loops reading
177/// opcode-prefixed tokens until that many line-end markers have been seen. Each
178/// line-end opcode increments the completed-line counter and pushes a
179/// [`Token::LineEnd`]; every other opcode is decoded into its token.
180///
181/// # Arguments
182///
183/// * `data` - The raw tokenized script bytes, starting with the line count.
184///
185/// # Returns
186///
187/// A [`TokenStream`] holding the declared line count and all parsed tokens.
188///
189/// # Errors
190///
191/// Returns [`TokenError::Truncated`] if the input ends before a required field
192/// can be read, or if the completed-line counter would overflow. Returns
193/// [`TokenError::BadString`] if a string token's length prefix or XOR-keyed
194/// UTF-16 code units are malformed.
195pub fn parse(data: &[u8]) -> Result<TokenStream, TokenError> {
196    let mut reader = TokenReader::new(data);
197    let line_count = reader.read_u32_le()?;
198    let mut completed_lines = 0u32;
199    let mut tokens = Vec::new();
200    while completed_lines < line_count {
201        let opcode = reader.read_u8()?;
202        if opcode == OPCODE_LINE_END {
203            completed_lines = completed_lines
204                .checked_add(1)
205                .ok_or(TokenError::Truncated)?;
206            tokens.push(Token::LineEnd);
207        } else {
208            tokens.push(read_token(opcode, &mut reader)?);
209        }
210    }
211    Ok(TokenStream { line_count, tokens })
212}
213
214/// Decodes a single non-`LineEnd` token from the given opcode.
215///
216/// Dispatches on `opcode`: id opcodes read an `i32` and resolve it through
217/// [`tables::keyword_by_id`] / [`tables::function_by_id`], falling back to the
218/// `Unknown*Id` variants; numeric opcodes read their fixed-width literal;
219/// string opcodes read an XOR-keyed UTF-16 payload (canonicalizing keyword,
220/// function, and macro casing). Any unrecognized opcode is mapped to an
221/// [`Token::Operator`] via [`operator`], or [`Token::UnknownOpcode`] if that
222/// also fails.
223///
224/// # Arguments
225///
226/// * `opcode` - The leading opcode byte that selects the token form.
227/// * `reader` - The [`TokenReader`] positioned just past the opcode.
228///
229/// # Returns
230///
231/// The decoded [`Token`].
232///
233/// # Errors
234///
235/// Returns [`TokenError::Truncated`] if the reader runs out of bytes for the
236/// token's payload, or [`TokenError::BadString`] if a string payload is
237/// malformed.
238fn read_token(opcode: u8, reader: &mut TokenReader<'_>) -> Result<Token, TokenError> {
239    match opcode {
240        OPCODE_KEYWORD_ID => {
241            let id = reader.read_i32_le()?;
242            Ok(tables::keyword_by_id(id).map_or(Token::UnknownKeywordId(id), Token::Keyword))
243        }
244        OPCODE_FUNCTION_ID => {
245            let id = reader.read_i32_le()?;
246            Ok(tables::function_by_id(id).map_or(Token::UnknownFunctionId(id), Token::Function))
247        }
248        OPCODE_U32 => Ok(Token::U32(reader.read_u32_le()?)),
249        OPCODE_U64 => Ok(Token::U64(reader.read_u64_le()?)),
250        OPCODE_F64 => Ok(Token::F64(reader.read_f64_le()?)),
251        OPCODE_KEYWORD_STRING => Ok(Token::Keyword(tables::canonical_keyword(
252            reader.read_xored_utf16_string()?.as_str(),
253        ))),
254        OPCODE_FUNCTION_STRING => Ok(Token::Function(tables::canonical_function(
255            reader.read_xored_utf16_string()?.as_str(),
256        ))),
257        OPCODE_MACRO_STRING => Ok(Token::Macro(tables::canonical_macro(
258            reader.read_xored_utf16_string()?.as_str(),
259        ))),
260        OPCODE_VARIABLE_STRING => Ok(Token::Variable(reader.read_xored_utf16_string()?)),
261        OPCODE_BARE_STRING => Ok(Token::BareString(reader.read_xored_utf16_string()?)),
262        OPCODE_PROPERTY_STRING => Ok(Token::Property(reader.read_xored_utf16_string()?)),
263        OPCODE_QUOTED_STRING => Ok(Token::QuotedString(reader.read_xored_utf16_string()?)),
264        OPCODE_RAW_STRING => Ok(Token::RawString(reader.read_xored_utf16_string()?)),
265        _ => operator(opcode).map_or(Ok(Token::UnknownOpcode(opcode)), |op| {
266            Ok(Token::Operator(op))
267        }),
268    }
269}
270
271/// Computes the indentation level to render the current line at.
272///
273/// Block-closing keywords (`Case`, `Else`, `ElseIf`, `WEnd`, `Until`, `Next`,
274/// `EndSelect`, `EndSwitch`, `EndFunc`, `EndIf`) are dedented one level
275/// relative to the surrounding block; every other line keeps the current
276/// `indent`. The dedent saturates at zero.
277///
278/// # Arguments
279///
280/// * `indent` - The current block indentation level.
281/// * `line` - The display strings of the tokens on the line being rendered.
282///
283/// # Returns
284///
285/// The number of leading tabs to emit for this line.
286fn line_indent(indent: usize, line: &[String]) -> usize {
287    match first_line_token(line) {
288        Some(
289            "Case" | "Else" | "ElseIf" | "WEnd" | "Until" | "Next" | "EndSelect" | "EndSwitch"
290            | "EndFunc" | "EndIf",
291        ) => indent.saturating_sub(1),
292        _ => indent,
293    }
294}
295
296/// Computes the indentation level for the line following the current one.
297///
298/// Block-opening keywords (`While`, `Do`, `For`, `Select`, `Switch`, `Func`)
299/// increase the level by one; block-closing keywords decrease it by one
300/// (saturating at zero). `If` is special: it opens a block only in its
301/// multi-line form, detected by `Then` being the last token on the line. A
302/// one-line `If <cond> Then <stmt>` has no matching `EndIf`, so it must not
303/// increase the indentation level.
304///
305/// # Arguments
306///
307/// * `indent` - The current block indentation level.
308/// * `line` - The display strings of the tokens on the line just rendered.
309///
310/// # Returns
311///
312/// The indentation level to apply to the next line.
313fn next_indent(indent: usize, line: &[String]) -> usize {
314    match first_line_token(line) {
315        // `If` opens a block only in its multi-line form, where `Then` is the
316        // final token on the line. A one-line `If <cond> Then <stmt>` has no
317        // matching `EndIf`, so it must not increase the indentation level.
318        Some("If") if last_line_token(line) == Some("Then") => indent.saturating_add(1),
319        Some("If") => indent,
320        Some("While" | "Do" | "For" | "Select" | "Switch" | "Func") => indent.saturating_add(1),
321        Some("WEnd" | "Until" | "Next" | "EndSelect" | "EndSwitch" | "EndFunc" | "EndIf") => {
322            indent.saturating_sub(1)
323        }
324        _ => indent,
325    }
326}
327
328/// Returns the first token's text on a line, used for indent decisions.
329///
330/// # Arguments
331///
332/// * `line` - The display strings of the tokens on a line.
333///
334/// # Returns
335///
336/// `Some` with the first token's text, or `None` if the line is empty.
337fn first_line_token(line: &[String]) -> Option<&str> {
338    line.first().map(String::as_str)
339}
340
341/// Returns the last token's text on a line, used to detect a trailing `Then`.
342///
343/// # Arguments
344///
345/// * `line` - The display strings of the tokens on a line.
346///
347/// # Returns
348///
349/// `Some` with the last token's text, or `None` if the line is empty.
350fn last_line_token(line: &[String]) -> Option<&str> {
351    line.last().map(String::as_str)
352}
353
354/// Joins a line's token display strings with spacing rules applied.
355///
356/// Tokens are concatenated with a single space between them, except where
357/// [`has_no_space_before`] suppresses the space before the current token or
358/// [`has_no_space_after`] suppresses it after the previous token (so commas
359/// hug the preceding token and brackets hug their operands).
360///
361/// # Arguments
362///
363/// * `line` - The display strings of the tokens on the line to render.
364///
365/// # Returns
366///
367/// The rendered line text without leading indentation or line terminator.
368fn render_line(line: &[String]) -> String {
369    let mut out = String::new();
370    let mut previous: Option<&str> = None;
371    for token in line {
372        let current = token.as_str();
373        if !out.is_empty()
374            && !has_no_space_before(current)
375            && !previous.is_some_and(has_no_space_after)
376        {
377            out.push(' ');
378        }
379        out.push_str(current);
380        previous = Some(current);
381    }
382    out
383}
384
385/// Reports whether no space should precede this token when rendering a line.
386///
387/// # Arguments
388///
389/// * `token` - The display text of the current token.
390///
391/// # Returns
392///
393/// `true` for `,`, `)`, `]`, `(`, and `[`, which hug the preceding token.
394fn has_no_space_before(token: &str) -> bool {
395    matches!(token, "," | ")" | "]" | "(" | "[")
396}
397
398/// Reports whether no space should follow this token when rendering a line.
399///
400/// # Arguments
401///
402/// * `token` - The display text of the previous token.
403///
404/// # Returns
405///
406/// `true` for `(` and `[`, so the following operand hugs the opening bracket.
407fn has_no_space_after(token: &str) -> bool {
408    matches!(token, "(" | "[")
409}
410
411/// Maps an operator/punctuation opcode to its source text.
412///
413/// Covers the contiguous `0x40`..=`0x58` opcode range, mapping each to its
414/// AutoIt operator or punctuation string (comparison, arithmetic, assignment,
415/// grouping, and the `?`/`:` ternary tokens).
416///
417/// # Arguments
418///
419/// * `opcode` - The opcode byte to translate.
420///
421/// # Returns
422///
423/// `Some` with the static operator text, or `None` if the opcode is not a
424/// known operator.
425fn operator(opcode: u8) -> Option<&'static str> {
426    match opcode {
427        0x40 => Some(","),
428        0x41 => Some("="),
429        0x42 => Some(">"),
430        0x43 => Some("<"),
431        0x44 => Some("<>"),
432        0x45 => Some(">="),
433        0x46 => Some("<="),
434        0x47 => Some("("),
435        0x48 => Some(")"),
436        0x49 => Some("+"),
437        0x4a => Some("-"),
438        0x4b => Some("/"),
439        0x4c => Some("*"),
440        0x4d => Some("&"),
441        0x4e => Some("["),
442        0x4f => Some("]"),
443        0x50 => Some("=="),
444        0x51 => Some("^"),
445        0x52 => Some("+="),
446        0x53 => Some("-="),
447        0x54 => Some("/="),
448        0x55 => Some("*="),
449        0x56 => Some("&="),
450        0x57 => Some("?"),
451        0x58 => Some(":"),
452        _ => None,
453    }
454}
455
456/// Cursor-tracking reader over the tokenized script bytes.
457struct TokenReader<'a> {
458    /// The full input byte slice being read.
459    data: &'a [u8],
460    /// The current read offset into `data`.
461    cursor: usize,
462}
463
464impl<'a> TokenReader<'a> {
465    /// Creates a reader positioned at the start of `data`.
466    ///
467    /// # Arguments
468    ///
469    /// * `data` - The byte slice to read tokens from.
470    ///
471    /// # Returns
472    ///
473    /// A [`TokenReader`] with its cursor at offset zero.
474    const fn new(data: &'a [u8]) -> Self {
475        Self { data, cursor: 0 }
476    }
477
478    /// Reads one byte and advances the cursor.
479    ///
480    /// # Returns
481    ///
482    /// The byte at the current cursor position.
483    ///
484    /// # Errors
485    ///
486    /// Returns [`TokenError::Truncated`] if the cursor is at or past the end of
487    /// the input, or if advancing the cursor would overflow.
488    fn read_u8(&mut self) -> Result<u8, TokenError> {
489        let byte = *self.data.get(self.cursor).ok_or(TokenError::Truncated)?;
490        self.cursor = self.cursor.checked_add(1).ok_or(TokenError::Truncated)?;
491        Ok(byte)
492    }
493
494    /// Reads a little-endian `u32` and advances the cursor by four bytes.
495    ///
496    /// # Returns
497    ///
498    /// The decoded unsigned 32-bit value.
499    ///
500    /// # Errors
501    ///
502    /// Returns [`TokenError::Truncated`] if fewer than four bytes remain, or if
503    /// advancing the cursor would overflow.
504    fn read_u32_le(&mut self) -> Result<u32, TokenError> {
505        let value = util::read_u32_le(self.data, self.cursor).ok_or(TokenError::Truncated)?;
506        self.cursor = self.cursor.checked_add(4).ok_or(TokenError::Truncated)?;
507        Ok(value)
508    }
509
510    /// Reads a little-endian `i32` by reinterpreting a `u32`'s bit pattern.
511    ///
512    /// # Returns
513    ///
514    /// The decoded signed 32-bit value.
515    ///
516    /// # Errors
517    ///
518    /// Returns [`TokenError::Truncated`] if fewer than four bytes remain.
519    fn read_i32_le(&mut self) -> Result<i32, TokenError> {
520        let value = self.read_u32_le()?;
521        Ok(i32::from_le_bytes(value.to_le_bytes()))
522    }
523
524    /// Reads a little-endian `u64` as two consecutive `u32` halves.
525    ///
526    /// The first `u32` supplies the low 32 bits and the second the high bits.
527    ///
528    /// # Returns
529    ///
530    /// The decoded unsigned 64-bit value.
531    ///
532    /// # Errors
533    ///
534    /// Returns [`TokenError::Truncated`] if fewer than eight bytes remain.
535    fn read_u64_le(&mut self) -> Result<u64, TokenError> {
536        let low = u64::from(self.read_u32_le()?);
537        let high = u64::from(self.read_u32_le()?);
538        Ok(low | (high << 32))
539    }
540
541    /// Reads a little-endian IEEE-754 `f64` from its raw 64-bit pattern.
542    ///
543    /// # Returns
544    ///
545    /// The decoded floating-point value.
546    ///
547    /// # Errors
548    ///
549    /// Returns [`TokenError::Truncated`] if fewer than eight bytes remain.
550    fn read_f64_le(&mut self) -> Result<f64, TokenError> {
551        Ok(f64::from_bits(self.read_u64_le()?))
552    }
553
554    /// Reads a length-prefixed, XOR-keyed UTF-16 string token.
555    ///
556    /// The leading `u32` is both the UTF-16 code-unit count and the XOR key:
557    /// each subsequent little-endian `u16` unit is XORed with the low 16 bits
558    /// of the key before being decoded. The recovered units are then assembled
559    /// into a [`String`].
560    ///
561    /// # Returns
562    ///
563    /// The decoded string.
564    ///
565    /// # Errors
566    ///
567    /// Returns [`TokenError::Truncated`] if the length prefix or any code unit
568    /// runs past the end of the input. Returns [`TokenError::BadString`] if the
569    /// key cannot be converted to the needed width or if the decoded units are
570    /// not valid UTF-16.
571    fn read_xored_utf16_string(&mut self) -> Result<String, TokenError> {
572        let key = self.read_u32_le()?;
573        let char_count = usize::try_from(key).map_err(|_err| TokenError::BadString)?;
574        // The length prefix is attacker-controlled and unbounded (up to
575        // `u32::MAX`), while every code unit consumes two bytes of input. Cap
576        // the up-front reservation at what the remaining input could actually
577        // supply, so a tiny malformed token cannot force a multi-gigabyte
578        // allocation; the loop below still fails with `Truncated` once the bytes
579        // run out. For a well-formed string the units are all present, so the
580        // bound never under-reserves. Mirrors the cap the JB decompressor applies.
581        let max_units = self.data.len().saturating_sub(self.cursor) / 2;
582        let mut units = Vec::with_capacity(char_count.min(max_units));
583        for _ in 0..char_count {
584            let raw = util::read_u16_le(self.data, self.cursor).ok_or(TokenError::Truncated)?;
585            self.cursor = self.cursor.checked_add(2).ok_or(TokenError::Truncated)?;
586            let decoded = raw ^ u16::try_from(key).map_err(|_err| TokenError::BadString)?;
587            units.push(decoded);
588        }
589        String::from_utf16(units.as_slice()).map_err(|_err| TokenError::BadString)
590    }
591}
592
593impl From<TokenError> for Error {
594    /// Converts a [`TokenError`] into the crate-level [`Error`].
595    ///
596    /// All token-parsing failures collapse into a single token error; the
597    /// specific [`TokenError`] variant is not preserved.
598    ///
599    /// # Arguments
600    ///
601    /// * `_value` - The [`TokenError`] to convert (its variant is discarded).
602    ///
603    /// # Returns
604    ///
605    /// The crate [`Error`] produced by `Error::token_error`.
606    fn from(_value: TokenError) -> Self {
607        Error::token_error()
608    }
609}
610
611#[cfg(test)]
612mod tests {
613    use super::*;
614
615    #[test]
616    fn parses_variable_assignment() -> Result<(), String> {
617        let mut data = Vec::new();
618        data.extend_from_slice(&1u32.to_le_bytes());
619        data.push(OPCODE_VARIABLE_STRING);
620        append_xored_string(&mut data, "x")?;
621        data.push(0x41);
622        data.push(OPCODE_U32);
623        data.extend_from_slice(&1u32.to_le_bytes());
624        data.push(OPCODE_LINE_END);
625
626        let stream = parse(data.as_slice()).map_err(|err| format!("{err:?}"))?;
627
628        check_eq(stream.line_count(), 1, "line count")?;
629        check_eq(stream.render_source(), "$x = 1\r\n".to_string(), "render")
630    }
631
632    #[test]
633    fn renders_simple_msgbox_call() -> Result<(), String> {
634        let mut data = Vec::new();
635        data.extend_from_slice(&1u32.to_le_bytes());
636        data.push(OPCODE_FUNCTION_ID);
637        data.extend_from_slice(&248i32.to_le_bytes());
638        data.push(0x47);
639        data.push(OPCODE_U32);
640        data.extend_from_slice(&0u32.to_le_bytes());
641        data.push(0x40);
642        data.push(OPCODE_QUOTED_STRING);
643        append_xored_string(&mut data, "title")?;
644        data.push(0x40);
645        data.push(OPCODE_QUOTED_STRING);
646        append_xored_string(&mut data, "text")?;
647        data.push(0x48);
648        data.push(OPCODE_LINE_END);
649
650        let stream = parse(data.as_slice()).map_err(|err| format!("{err:?}"))?;
651
652        check_eq(
653            stream.tokens(),
654            &[
655                Token::Function("MsgBox".to_string()),
656                Token::Operator("("),
657                Token::U32(0),
658                Token::Operator(","),
659                Token::QuotedString("title".to_string()),
660                Token::Operator(","),
661                Token::QuotedString("text".to_string()),
662                Token::Operator(")"),
663                Token::LineEnd,
664            ],
665            "tokens",
666        )?;
667        check_eq(
668            stream.render_source(),
669            "MsgBox(0, \"title\", \"text\")\r\n".to_string(),
670            "render",
671        )
672    }
673
674    #[test]
675    fn preserves_unknown_opcode() -> Result<(), String> {
676        let data = [1, 0, 0, 0, 0xff, OPCODE_LINE_END];
677        let stream = parse(&data).map_err(|err| format!("{err:?}"))?;
678
679        check_eq(
680            stream.tokens(),
681            &[Token::UnknownOpcode(0xff), Token::LineEnd],
682            "tokens",
683        )
684    }
685
686    #[test]
687    fn oversized_string_length_fails_fast_without_overallocating() -> Result<(), String> {
688        // A string token whose length prefix claims `u32::MAX` code units but is
689        // backed by no actual units must fail fast with `Truncated` rather than
690        // reserving gigabytes up front. The reservation is bounded by the
691        // remaining input (zero here), so the read loop runs out of bytes on its
692        // first iteration instead of allocating ~8 GiB.
693        let mut data = Vec::new();
694        data.extend_from_slice(&1u32.to_le_bytes());
695        data.push(OPCODE_VARIABLE_STRING);
696        data.extend_from_slice(&u32::MAX.to_le_bytes());
697
698        match parse(data.as_slice()) {
699            Err(TokenError::Truncated) => Ok(()),
700            other => Err(format!("expected Truncated, got {other:?}")),
701        }
702    }
703
704    #[test]
705    fn parses_quoted_macro_and_unknown_ids() -> Result<(), String> {
706        let mut data = Vec::new();
707        data.extend_from_slice(&1u32.to_le_bytes());
708        data.push(OPCODE_KEYWORD_ID);
709        data.extend_from_slice(&4i32.to_le_bytes());
710        data.push(OPCODE_MACRO_STRING);
711        append_xored_string(&mut data, "ScriptName")?;
712        data.push(OPCODE_QUOTED_STRING);
713        append_xored_string(&mut data, "a\"b")?;
714        data.push(OPCODE_KEYWORD_ID);
715        data.extend_from_slice(&999i32.to_le_bytes());
716        data.push(OPCODE_FUNCTION_ID);
717        data.extend_from_slice(&999i32.to_le_bytes());
718        data.push(OPCODE_LINE_END);
719
720        let stream = parse(data.as_slice()).map_err(|err| format!("{err:?}"))?;
721
722        check_eq(
723            stream.tokens(),
724            &[
725                Token::Keyword("If".to_string()),
726                Token::Macro("ScriptName".to_string()),
727                Token::QuotedString("a\"b".to_string()),
728                Token::UnknownKeywordId(999),
729                Token::UnknownFunctionId(999),
730                Token::LineEnd,
731            ],
732            "tokens",
733        )?;
734        check_eq(
735            stream.render_source(),
736            "If @ScriptName \"a\"\"b\" <keyword:999> <function:999>\r\n".to_string(),
737            "render",
738        )
739    }
740
741    #[test]
742    fn resolves_function_ids_and_canonicalizes_strings() -> Result<(), String> {
743        let mut data = Vec::new();
744        data.extend_from_slice(&1u32.to_le_bytes());
745        data.push(OPCODE_FUNCTION_ID);
746        data.extend_from_slice(&248i32.to_le_bytes());
747        data.push(OPCODE_FUNCTION_STRING);
748        append_xored_string(&mut data, "runwait")?;
749        data.push(OPCODE_MACRO_STRING);
750        append_xored_string(&mut data, "scriptname")?;
751        data.push(OPCODE_LINE_END);
752
753        let stream = parse(data.as_slice()).map_err(|err| format!("{err:?}"))?;
754
755        check_eq(
756            stream.tokens(),
757            &[
758                Token::Function("MsgBox".to_string()),
759                Token::Function("RunWait".to_string()),
760                Token::Macro("ScriptName".to_string()),
761                Token::LineEnd,
762            ],
763            "tokens",
764        )
765    }
766
767    #[test]
768    fn resolves_expanded_low_function_ids() -> Result<(), String> {
769        let mut data = Vec::new();
770        data.extend_from_slice(&1u32.to_le_bytes());
771        data.push(OPCODE_FUNCTION_ID);
772        data.extend_from_slice(&12i32.to_le_bytes());
773        data.push(OPCODE_FUNCTION_ID);
774        data.extend_from_slice(&17i32.to_le_bytes());
775        data.push(OPCODE_FUNCTION_ID);
776        data.extend_from_slice(&27i32.to_le_bytes());
777        data.push(OPCODE_FUNCTION_STRING);
778        append_xored_string(&mut data, "binarytostring")?;
779        data.push(OPCODE_LINE_END);
780
781        let stream = parse(data.as_slice()).map_err(|err| format!("{err:?}"))?;
782
783        check_eq(
784            stream.tokens(),
785            &[
786                Token::Function("Beep".to_string()),
787                Token::Function("BitAND".to_string()),
788                Token::Function("Ceiling".to_string()),
789                Token::Function("BinaryToString".to_string()),
790                Token::LineEnd,
791            ],
792            "tokens",
793        )
794    }
795
796    #[test]
797    fn resolves_expanded_control_and_directory_function_ids() -> Result<(), String> {
798        let mut data = Vec::new();
799        data.extend_from_slice(&1u32.to_le_bytes());
800        data.push(OPCODE_FUNCTION_ID);
801        data.extend_from_slice(&30i32.to_le_bytes());
802        data.push(OPCODE_FUNCTION_ID);
803        data.extend_from_slice(&45i32.to_le_bytes());
804        data.push(OPCODE_FUNCTION_ID);
805        data.extend_from_slice(&56i32.to_le_bytes());
806        data.push(OPCODE_FUNCTION_STRING);
807        append_xored_string(&mut data, "consolewriteerror")?;
808        data.push(OPCODE_LINE_END);
809
810        let stream = parse(data.as_slice()).map_err(|err| format!("{err:?}"))?;
811
812        check_eq(
813            stream.tokens(),
814            &[
815                Token::Function("ClipGet".to_string()),
816                Token::Function("ControlListView".to_string()),
817                Token::Function("DirMove".to_string()),
818                Token::Function("ConsoleWriteError".to_string()),
819                Token::LineEnd,
820            ],
821            "tokens",
822        )
823    }
824
825    #[test]
826    fn resolves_expanded_dll_drive_and_env_function_ids() -> Result<(), String> {
827        let mut data = Vec::new();
828        data.extend_from_slice(&1u32.to_le_bytes());
829        data.push(OPCODE_FUNCTION_ID);
830        data.extend_from_slice(&59i32.to_le_bytes());
831        data.push(OPCODE_FUNCTION_ID);
832        data.extend_from_slice(&68i32.to_le_bytes());
833        data.push(OPCODE_FUNCTION_ID);
834        data.extend_from_slice(&84i32.to_le_bytes());
835        data.push(OPCODE_FUNCTION_ID);
836        data.extend_from_slice(&85i32.to_le_bytes());
837        data.push(OPCODE_FUNCTION_ID);
838        data.extend_from_slice(&86i32.to_le_bytes());
839        data.push(OPCODE_FUNCTION_STRING);
840        append_xored_string(&mut data, "drivegetfilesystem")?;
841        data.push(OPCODE_LINE_END);
842
843        let stream = parse(data.as_slice()).map_err(|err| format!("{err:?}"))?;
844
845        check_eq(
846            stream.tokens(),
847            &[
848                Token::Function("DllCallbackFree".to_string()),
849                Token::Function("DllStructSetData".to_string()),
850                Token::Function("EnvUpdate".to_string()),
851                Token::Function("Eval".to_string()),
852                Token::Function("Execute".to_string()),
853                Token::Function("DriveGetFileSystem".to_string()),
854                Token::LineEnd,
855            ],
856            "tokens",
857        )
858    }
859
860    #[test]
861    fn resolves_expanded_file_setup_and_metadata_function_ids() -> Result<(), String> {
862        let mut data = Vec::new();
863        data.extend_from_slice(&1u32.to_le_bytes());
864        data.push(OPCODE_FUNCTION_ID);
865        data.extend_from_slice(&87i32.to_le_bytes());
866        data.push(OPCODE_FUNCTION_ID);
867        data.extend_from_slice(&91i32.to_le_bytes());
868        data.push(OPCODE_FUNCTION_ID);
869        data.extend_from_slice(&95i32.to_le_bytes());
870        data.push(OPCODE_FUNCTION_ID);
871        data.extend_from_slice(&99i32.to_le_bytes());
872        data.push(OPCODE_FUNCTION_ID);
873        data.extend_from_slice(&106i32.to_le_bytes());
874        data.push(OPCODE_FUNCTION_STRING);
875        append_xored_string(&mut data, "filegetshortcut")?;
876        data.push(OPCODE_LINE_END);
877
878        let stream = parse(data.as_slice()).map_err(|err| format!("{err:?}"))?;
879
880        check_eq(
881            stream.tokens(),
882            &[
883                Token::Function("Exp".to_string()),
884                Token::Function("FileCreateNTFSLink".to_string()),
885                Token::Function("FileFindFirstFile".to_string()),
886                Token::Function("FileGetEncoding".to_string()),
887                Token::Function("FileGetVersion".to_string()),
888                Token::Function("FileGetShortcut".to_string()),
889                Token::LineEnd,
890            ],
891            "tokens",
892        )
893    }
894
895    #[test]
896    fn resolves_expanded_file_io_function_ids() -> Result<(), String> {
897        let mut data = Vec::new();
898        data.extend_from_slice(&1u32.to_le_bytes());
899        data.push(OPCODE_FUNCTION_ID);
900        data.extend_from_slice(&107i32.to_le_bytes());
901        data.push(OPCODE_FUNCTION_ID);
902        data.extend_from_slice(&111i32.to_le_bytes());
903        data.push(OPCODE_FUNCTION_ID);
904        data.extend_from_slice(&113i32.to_le_bytes());
905        data.push(OPCODE_FUNCTION_ID);
906        data.extend_from_slice(&123i32.to_le_bytes());
907        data.push(OPCODE_FUNCTION_ID);
908        data.extend_from_slice(&126i32.to_le_bytes());
909        data.push(OPCODE_FUNCTION_STRING);
910        append_xored_string(&mut data, "filesavedialog")?;
911        data.push(OPCODE_LINE_END);
912
913        let stream = parse(data.as_slice()).map_err(|err| format!("{err:?}"))?;
914
915        check_eq(
916            stream.tokens(),
917            &[
918                Token::Function("FileInstall".to_string()),
919                Token::Function("FileRead".to_string()),
920                Token::Function("FileReadToArray".to_string()),
921                Token::Function("FileWriteLine".to_string()),
922                Token::Function("FuncName".to_string()),
923                Token::Function("FileSaveDialog".to_string()),
924                Token::LineEnd,
925            ],
926            "tokens",
927        )
928    }
929
930    #[test]
931    fn resolves_expanded_gui_http_and_inet_function_ids() -> Result<(), String> {
932        let mut data = Vec::new();
933        data.extend_from_slice(&1u32.to_le_bytes());
934        data.push(OPCODE_FUNCTION_ID);
935        data.extend_from_slice(&127i32.to_le_bytes());
936        data.push(OPCODE_FUNCTION_ID);
937        data.extend_from_slice(&142i32.to_le_bytes());
938        data.push(OPCODE_FUNCTION_ID);
939        data.extend_from_slice(&162i32.to_le_bytes());
940        data.push(OPCODE_FUNCTION_ID);
941        data.extend_from_slice(&181i32.to_le_bytes());
942        data.push(OPCODE_FUNCTION_ID);
943        data.extend_from_slice(&200i32.to_le_bytes());
944        data.push(OPCODE_FUNCTION_ID);
945        data.extend_from_slice(&204i32.to_le_bytes());
946        data.push(OPCODE_FUNCTION_ID);
947        data.extend_from_slice(&207i32.to_le_bytes());
948        data.push(OPCODE_FUNCTION_STRING);
949        append_xored_string(&mut data, "guictrlsetbkcolor")?;
950        data.push(OPCODE_LINE_END);
951
952        let stream = parse(data.as_slice()).map_err(|err| format!("{err:?}"))?;
953
954        check_eq(
955            stream.tokens(),
956            &[
957                Token::Function("GUICreate".to_string()),
958                Token::Function("GUICtrlCreateListView".to_string()),
959                Token::Function("GUICtrlRegisterListViewSort".to_string()),
960                Token::Function("GUIDelete".to_string()),
961                Token::Function("HttpSetProxy".to_string()),
962                Token::Function("InetGet".to_string()),
963                Token::Function("InetRead".to_string()),
964                Token::Function("GUICtrlSetBkColor".to_string()),
965                Token::LineEnd,
966            ],
967            "tokens",
968        )
969    }
970
971    #[test]
972    fn resolves_expanded_ini_type_map_mouse_and_msgbox_function_ids() -> Result<(), String> {
973        let mut data = Vec::new();
974        data.extend_from_slice(&1u32.to_le_bytes());
975        data.push(OPCODE_FUNCTION_ID);
976        data.extend_from_slice(&208i32.to_le_bytes());
977        data.push(OPCODE_FUNCTION_ID);
978        data.extend_from_slice(&211i32.to_le_bytes());
979        data.push(OPCODE_FUNCTION_ID);
980        data.extend_from_slice(&228i32.to_le_bytes());
981        data.push(OPCODE_FUNCTION_ID);
982        data.extend_from_slice(&235i32.to_le_bytes());
983        data.push(OPCODE_FUNCTION_ID);
984        data.extend_from_slice(&241i32.to_le_bytes());
985        data.push(OPCODE_FUNCTION_ID);
986        data.extend_from_slice(&248i32.to_le_bytes());
987        data.push(OPCODE_FUNCTION_STRING);
988        append_xored_string(&mut data, "isstring")?;
989        data.push(OPCODE_LINE_END);
990
991        let stream = parse(data.as_slice()).map_err(|err| format!("{err:?}"))?;
992
993        check_eq(
994            stream.tokens(),
995            &[
996                Token::Function("IniDelete".to_string()),
997                Token::Function("IniReadSectionNames".to_string()),
998                Token::Function("IsMap".to_string()),
999                Token::Function("MapExists".to_string()),
1000                Token::Function("MouseClickDrag".to_string()),
1001                Token::Function("MsgBox".to_string()),
1002                Token::Function("IsString".to_string()),
1003                Token::LineEnd,
1004            ],
1005            "tokens",
1006        )
1007    }
1008
1009    #[test]
1010    fn resolves_expanded_object_process_registry_and_run_function_ids() -> Result<(), String> {
1011        let mut data = Vec::new();
1012        data.extend_from_slice(&1u32.to_le_bytes());
1013        data.push(OPCODE_FUNCTION_ID);
1014        data.extend_from_slice(&249i32.to_le_bytes());
1015        data.push(OPCODE_FUNCTION_ID);
1016        data.extend_from_slice(&255i32.to_le_bytes());
1017        data.push(OPCODE_FUNCTION_ID);
1018        data.extend_from_slice(&263i32.to_le_bytes());
1019        data.push(OPCODE_FUNCTION_ID);
1020        data.extend_from_slice(&277i32.to_le_bytes());
1021        data.push(OPCODE_FUNCTION_ID);
1022        data.extend_from_slice(&280i32.to_le_bytes());
1023        data.push(OPCODE_FUNCTION_ID);
1024        data.extend_from_slice(&288i32.to_le_bytes());
1025        data.push(OPCODE_FUNCTION_ID);
1026        data.extend_from_slice(&300i32.to_le_bytes());
1027        data.push(OPCODE_FUNCTION_STRING);
1028        append_xored_string(&mut data, "soundsetwavevolume")?;
1029        data.push(OPCODE_LINE_END);
1030
1031        let stream = parse(data.as_slice()).map_err(|err| format!("{err:?}"))?;
1032
1033        check_eq(
1034            stream.tokens(),
1035            &[
1036                Token::Function("Number".to_string()),
1037                Token::Function("OnAutoItExitRegister".to_string()),
1038                Token::Function("ProcessExists".to_string()),
1039                Token::Function("RegRead".to_string()),
1040                Token::Function("Run".to_string()),
1041                Token::Function("ShellExecute".to_string()),
1042                Token::Function("StatusbarGetText".to_string()),
1043                Token::Function("SoundSetWaveVolume".to_string()),
1044                Token::LineEnd,
1045            ],
1046            "tokens",
1047        )
1048    }
1049
1050    #[test]
1051    fn resolves_final_string_network_tray_and_window_function_ids() -> Result<(), String> {
1052        let mut data = Vec::new();
1053        data.extend_from_slice(&1u32.to_le_bytes());
1054        data.push(OPCODE_FUNCTION_ID);
1055        data.extend_from_slice(&301i32.to_le_bytes());
1056        data.push(OPCODE_FUNCTION_ID);
1057        data.extend_from_slice(&310i32.to_le_bytes());
1058        data.push(OPCODE_FUNCTION_ID);
1059        data.extend_from_slice(&325i32.to_le_bytes());
1060        data.push(OPCODE_FUNCTION_ID);
1061        data.extend_from_slice(&339i32.to_le_bytes());
1062        data.push(OPCODE_FUNCTION_ID);
1063        data.extend_from_slice(&368i32.to_le_bytes());
1064        data.push(OPCODE_FUNCTION_ID);
1065        data.extend_from_slice(&393i32.to_le_bytes());
1066        data.push(OPCODE_FUNCTION_ID);
1067        data.extend_from_slice(&404i32.to_le_bytes());
1068        data.push(OPCODE_FUNCTION_STRING);
1069        append_xored_string(&mut data, "stringtoasciiarray")?;
1070        data.push(OPCODE_LINE_END);
1071
1072        let stream = parse(data.as_slice()).map_err(|err| format!("{err:?}"))?;
1073
1074        check_eq(
1075            stream.tokens(),
1076            &[
1077                Token::Function("StdErrRead".to_string()),
1078                Token::Function("StringInStr".to_string()),
1079                Token::Function("StringRegExp".to_string()),
1080                Token::Function("TCPAccept".to_string()),
1081                Token::Function("UBound".to_string()),
1082                Token::Function("WinMenuSelectItem".to_string()),
1083                Token::Function("WinWaitNotActive".to_string()),
1084                Token::Function("StringToASCIIArray".to_string()),
1085                Token::LineEnd,
1086            ],
1087            "tokens",
1088        )
1089    }
1090
1091    #[test]
1092    fn canonicalizes_expanded_official_macros() -> Result<(), String> {
1093        let mut data = Vec::new();
1094        data.extend_from_slice(&1u32.to_le_bytes());
1095        data.push(OPCODE_MACRO_STRING);
1096        append_xored_string(&mut data, "appdatacommondir")?;
1097        data.push(OPCODE_MACRO_STRING);
1098        append_xored_string(&mut data, "gui_ctrlhandle")?;
1099        data.push(OPCODE_MACRO_STRING);
1100        append_xored_string(&mut data, "sw_shownoactivate")?;
1101        data.push(OPCODE_MACRO_STRING);
1102        append_xored_string(&mut data, "tray_id")?;
1103        data.push(OPCODE_MACRO_STRING);
1104        append_xored_string(&mut data, "year")?;
1105        data.push(OPCODE_LINE_END);
1106
1107        let stream = parse(data.as_slice()).map_err(|err| format!("{err:?}"))?;
1108
1109        check_eq(
1110            stream.tokens(),
1111            &[
1112                Token::Macro("AppDataCommonDir".to_string()),
1113                Token::Macro("GUI_CtrlHandle".to_string()),
1114                Token::Macro("SW_SHOWNOACTIVATE".to_string()),
1115                Token::Macro("TRAY_ID".to_string()),
1116                Token::Macro("YEAR".to_string()),
1117                Token::LineEnd,
1118            ],
1119            "tokens",
1120        )
1121    }
1122
1123    #[test]
1124    fn renders_control_flow_with_indentation() -> Result<(), String> {
1125        let mut data = Vec::new();
1126        data.extend_from_slice(&3u32.to_le_bytes());
1127        data.push(OPCODE_KEYWORD_ID);
1128        data.extend_from_slice(&4i32.to_le_bytes());
1129        data.push(OPCODE_VARIABLE_STRING);
1130        append_xored_string(&mut data, "x")?;
1131        data.push(OPCODE_KEYWORD_ID);
1132        data.extend_from_slice(&5i32.to_le_bytes());
1133        data.push(OPCODE_LINE_END);
1134        data.push(OPCODE_VARIABLE_STRING);
1135        append_xored_string(&mut data, "x")?;
1136        data.push(0x41);
1137        data.push(OPCODE_U32);
1138        data.extend_from_slice(&1u32.to_le_bytes());
1139        data.push(OPCODE_LINE_END);
1140        data.push(OPCODE_KEYWORD_ID);
1141        data.extend_from_slice(&8i32.to_le_bytes());
1142        data.push(OPCODE_LINE_END);
1143
1144        let stream = parse(data.as_slice()).map_err(|err| format!("{err:?}"))?;
1145
1146        check_eq(
1147            stream.render_source(),
1148            "If $x Then\r\n\t$x = 1\r\nEndIf\r\n".to_string(),
1149            "render",
1150        )
1151    }
1152
1153    #[test]
1154    fn one_line_if_does_not_increase_indentation() -> Result<(), String> {
1155        // `If $x Then $x = 1` is a one-line statement with no `EndIf`; the
1156        // following line must stay at the outer indentation level rather than
1157        // accumulating a phantom block indent.
1158        let mut data = Vec::new();
1159        data.extend_from_slice(&2u32.to_le_bytes());
1160        // If $x Then $x = 1
1161        data.push(OPCODE_KEYWORD_ID);
1162        data.extend_from_slice(&4i32.to_le_bytes());
1163        data.push(OPCODE_VARIABLE_STRING);
1164        append_xored_string(&mut data, "x")?;
1165        data.push(OPCODE_KEYWORD_ID);
1166        data.extend_from_slice(&5i32.to_le_bytes());
1167        data.push(OPCODE_VARIABLE_STRING);
1168        append_xored_string(&mut data, "x")?;
1169        data.push(0x41);
1170        data.push(OPCODE_U32);
1171        data.extend_from_slice(&1u32.to_le_bytes());
1172        data.push(OPCODE_LINE_END);
1173        // $y = 2
1174        data.push(OPCODE_VARIABLE_STRING);
1175        append_xored_string(&mut data, "y")?;
1176        data.push(0x41);
1177        data.push(OPCODE_U32);
1178        data.extend_from_slice(&2u32.to_le_bytes());
1179        data.push(OPCODE_LINE_END);
1180
1181        let stream = parse(data.as_slice()).map_err(|err| format!("{err:?}"))?;
1182
1183        check_eq(
1184            stream.render_source(),
1185            "If $x Then $x = 1\r\n$y = 2\r\n".to_string(),
1186            "render",
1187        )
1188    }
1189
1190    fn append_xored_string(out: &mut Vec<u8>, value: &str) -> Result<(), String> {
1191        let units: Vec<u16> = value.encode_utf16().collect();
1192        let key = u32::try_from(units.len()).map_err(|err| err.to_string())?;
1193        out.extend_from_slice(&key.to_le_bytes());
1194        let key16 = u16::try_from(key).map_err(|err| err.to_string())?;
1195        for unit in units {
1196            out.extend_from_slice(&(unit ^ key16).to_le_bytes());
1197        }
1198        Ok(())
1199    }
1200
1201    fn check_eq<T>(actual: T, expected: T, context: &str) -> Result<(), String>
1202    where
1203        T: core::fmt::Debug + PartialEq,
1204    {
1205        if actual == expected {
1206            Ok(())
1207        } else {
1208            Err(format!("{context}: got {actual:?}, expected {expected:?}"))
1209        }
1210    }
1211}