Repository: awslabs/llrt
File: modules/llrt_fetch/src/headers.rs (Rust)

Code before:
        if c == '\r' {
            if chars.peek() == Some(&'\n') {
                chars.next();
                match chars.peek() {
                    Some(' ') | Some('\t') => {
                        chars.next();
                        normalized.push(' ');
                        continue;
                    },
                    _ => continue,
                }
            } else {
                continue;
            }
        } else if c == '\n' || c == '\0' || c == '\u{000B}'
        /* \f */
        {
            continue;
        } else {
            normalized.push(c);
        }
    }

    let trimmed = normalized.trim_matches(|c| c == ' ' || c == '\t');
    Some(trimmed.to_string())
}

#[cfg(test)]
mod tests {
    use llrt_test::test_async_with;

    use super::*;

    #[tokio::test]
    async fn test_get_header() {
        test_async_with(|ctx| {
            crate::init(&ctx).unwrap();
            Box::pin(async move {
                let mut headers = Headers::new(ctx.clone(), Opt(None)).unwrap();
                headers
                    .set(
                        ctx.clone(),
                        "Content-Type".into(),
                        "application/json".into(),
                    )
                    .unwrap();
                headers
                    .append(ctx.clone(), "set-cookie".into(), "cookie1=value1".into())
                    .unwrap();
                headers
                    .append(ctx.clone(), "set-cookie".into(), "cookie2=value2".into())

Review comment (suggestion):
This looks strange to me. We're returning an option when we always have `Some(value)`. We can also take ownership of the string and do inplace replacements. Example inplace modification:
```
fn urlify_spaces(text: &mut String) {
    const SPACE_REPLACEMENT: &[u8] = b"%20";

    // operating on bytes for simplicity
    let mut buffer = std::mem::take(text).into_bytes();
    let old_len = buffer.len();

    let space_count = buffer.iter().filter(|&&byte| byte == b' ').count();
    let new_len = buffer.len() + (SPACE_REPLACEMENT.len() - 1) * space_count;
    buffer.resize(new_len, b'\0');

    let mut write_pos = new_len;

    for read_pos in (0..old_len).rev() {
        let byte = buffer[read_pos];

        if byte == b' ' {
            write_pos -= SPACE_REPLACEMENT.len();
            buffer[write_pos..write_pos + SPACE_REPLACEMENT.len()]
                .copy_from_slice(SPACE_REPLACEMENT);
        } else {
            write_pos -= 1;
            buffer[write_pos] = byte;
        }
    }

    *text = String::from_utf8(buffer).expect("invalid UTF-8 during URL-ification");
}
```

Code after:
    // Skip leading SP or HTAB
    while read_idx < input.len() && (input[read_idx] == b' ' || input[read_idx] == b'\t') {
        read_idx += 1;
    }

    // Store the last whitespace byte if any (space or tab)
    let mut pending_whitespace: Option<u8> = None;

    while read_idx < input.len() {
        match input[read_idx] {
            // obs-fold: CRLF followed by SP or HTAB
            b'\r'
                if read_idx + 2 < input.len()
                    && input[read_idx + 1] == b'\n'
                    && (input[read_idx + 2] == b' ' || input[read_idx + 2] == b'\t') =>
            {
                pending_whitespace = Some(input[read_idx + 2]);
                read_idx += 3;
            },
            b'\r' | b'\n' => {
                // skip bare CR or LF
                read_idx += 1;
            },
            b' ' | b'\t' => {
                // keep the last whitespace char to write later (collapse multiple)
                pending_whitespace = Some(input[read_idx]);
                read_idx += 1;
            },
            byte => {
                // write pending whitespace if any
                if let Some(ws) = pending_whitespace.take() {
                    if write_idx > 0 {
                        input[write_idx] = ws;
                        write_idx += 1;
                    }
                }
                input[write_idx] = byte;
                write_idx += 1;
                read_idx += 1;
            },
        }
    }

    // Trim trailing SP or HTAB
    while write_idx > 0 && (input[write_idx - 1] == b' ' || input[write_idx - 1] == b'\t') {
        write_idx -= 1;
    }

    input.truncate(write_idx);
    *text = String::from_utf8(input).or_throw(ctx)?;
    Ok(())