ladybird/Userland/Libraries/LibIMAP/Parser.cpp

/*
 * Copyright (c) 2021, Kyle Pereira <hey@xylepereira.me>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include <AK/CharacterTypes.h>
#include <LibIMAP/Parser.h>

namespace IMAP {

ParseStatus Parser::parse(ByteBuffer&& buffer, bool expecting_tag)
{
    if (m_incomplete) {
        m_buffer += buffer;
        m_incomplete = false;
    } else {
        m_buffer = move(buffer);
        position = 0;
        m_response = SolidResponse();
    }

    if (try_consume("+")) {
        consume(" ");
        auto data = parse_while([](u8 x) { return x != '\r'; });
        consume("\r\n");
        return { true, { ContinueRequest { data } } };
    }

    while (try_consume("*")) {
        parse_untagged();
    }

    if (expecting_tag) {
        if (at_end()) {
            m_incomplete = true;
            return { true, {} };
        }
        parse_response_done();
    }

    if (m_parsing_failed) {
        return { false, {} };
    } else {
        return { true, { { move(m_response) } } };
    }
}

bool Parser::try_consume(StringView x)
{
    size_t i = 0;
    auto previous_position = position;
    while (i < x.length() && !at_end() && to_ascii_lowercase(x[i]) == to_ascii_lowercase(m_buffer[position])) {
        i++;
        position++;
    }
    if (i != x.length()) {
        // We didn't match the full string.
        position = previous_position;
        return false;
    }

    return true;
}

void Parser::parse_response_done()
{
    consume("A");
    auto tag = parse_number();
    consume(" ");

    ResponseStatus status = parse_status();
    consume(" ");

    m_response.m_tag = tag;
    m_response.m_status = status;

    StringBuilder response_data;

    while (!at_end() && m_buffer[position] != '\r') {
        response_data.append((char)m_buffer[position]);
        position += 1;
    }

    consume("\r\n");
    m_response.m_response_text = response_data.build();
}

void Parser::consume(StringView x)
{
    if (!try_consume(x)) {
        dbgln("{} not matched at {}, buffer: {}", x, position, StringView(m_buffer.data(), m_buffer.size()));

        m_parsing_failed = true;
    }
}

Optional<unsigned> Parser::try_parse_number()
{
    auto number_matched = 0;
    while (!at_end() && 0 <= m_buffer[position] - '0' && m_buffer[position] - '0' <= 9) {
        number_matched++;
        position++;
    }
    if (number_matched == 0)
        return {};

    auto number = StringView(m_buffer.data() + position - number_matched, number_matched);

    return number.to_uint();
}

unsigned Parser::parse_number()
{
    auto number = try_parse_number();
    if (!number.has_value()) {
        m_parsing_failed = true;
        return -1;
    }

    return number.value();
}

void Parser::parse_untagged()
{
    consume(" ");

    // Certain messages begin with a number like:
    // * 15 EXISTS
    auto number = try_parse_number();
    if (number.has_value()) {
        consume(" ");
        auto data_type = parse_atom().to_string();
        if (data_type.matches("EXISTS")) {
            m_response.data().set_exists(number.value());
            consume("\r\n");
        } else if (data_type.matches("RECENT")) {
            m_response.data().set_recent(number.value());
            consume("\r\n");
        } else if (data_type.matches("FETCH")) {
            auto fetch_response = parse_fetch_response();
            m_response.data().add_fetch_response(number.value(), move(fetch_response));
        }
        return;
    }

    if (try_consume("CAPABILITY")) {
        parse_capability_response();
    } else if (try_consume("LIST")) {
        auto item = parse_list_item();
        m_response.data().add_list_item(move(item));
    } else if (try_consume("FLAGS")) {
        consume(" ");
        auto flags = parse_list(+[](StringView x) { return String(x); });
        m_response.data().set_flags(move(flags));
        consume("\r\n");
    } else if (try_consume("OK")) {
        consume(" ");
        if (try_consume("[")) {
            auto actual_type = parse_atom();
            consume(" ");
            if (actual_type.matches("UIDNEXT")) {
                auto n = parse_number();
                m_response.data().set_uid_next(n);
            } else if (actual_type.matches("UIDVALIDITY")) {
                auto n = parse_number();
                m_response.data().set_uid_validity(n);
            } else if (actual_type.matches("UNSEEN")) {
                auto n = parse_number();
                m_response.data().set_unseen(n);
            } else if (actual_type.matches("PERMANENTFLAGS")) {
                auto flags = parse_list(+[](StringView x) { return String(x); });
                m_response.data().set_permanent_flags(move(flags));
            } else {
                dbgln("Unknown: {}", actual_type);
                parse_while([](u8 x) { return x != ']'; });
            }
            consume("]");
            parse_while([](u8 x) { return x != '\r'; });
            consume("\r\n");
        } else {
            parse_while([](u8 x) { return x != '\r'; });
            consume("\r\n");
        }
    } else if (try_consume("BYE")) {
        auto message = parse_while([](u8 x) { return x != '\r'; });
        consume("\r\n");
        m_response.data().set_bye(message.is_empty() ? Optional<String>() : Optional<String>(message));
    } else {
        auto x = parse_while([](u8 x) { return x != '\r'; });
        consume("\r\n");
        dbgln("ignored {}", x);
    }
}

StringView Parser::parse_quoted_string()
{
    auto str = parse_while([](u8 x) { return x != '"'; });
    consume("\"");
    return str;
}

StringView Parser::parse_string()
{
    if (try_consume("\"")) {
        return parse_quoted_string();
    } else {
        return parse_literal_string();
    }
}

Optional<StringView> Parser::parse_nstring()
{
    if (try_consume("NIL"))
        return {};
    else
        return { parse_string() };
}

FetchResponseData Parser::parse_fetch_response()
{
    consume(" (");
    auto fetch_response = FetchResponseData();

    while (!try_consume(")")) {
        auto data_item = parse_fetch_data_item();
        switch (data_item.type) {
        case FetchCommand::DataItemType::Envelope: {
            consume(" (");
            auto date = parse_nstring();
            consume(" ");
            auto subject = parse_nstring();
            consume(" ");
            auto from = parse_address_list();
            consume(" ");
            auto sender = parse_address_list();
            consume(" ");
            auto reply_to = parse_address_list();
            consume(" ");
            auto to = parse_address_list();
            consume(" ");
            auto cc = parse_address_list();
            consume(" ");
            auto bcc = parse_address_list();
            consume(" ");
            auto in_reply_to = parse_nstring();
            consume(" ");
            auto message_id = parse_nstring();
            consume(")");
            Envelope envelope = {
                date.has_value() ? Optional<String>(date.value()) : Optional<String>(),
                subject.has_value() ? Optional<String>(subject.value()) : Optional<String>(),
                from,
                sender,
                reply_to,
                to,
                cc,
                bcc,
                in_reply_to.has_value() ? Optional<String>(in_reply_to.value()) : Optional<String>(),
                message_id.has_value() ? Optional<String>(message_id.value()) : Optional<String>(),
            };
            fetch_response.set_envelope(move(envelope));
            break;
        }
        case FetchCommand::DataItemType::Flags: {
            consume(" ");
            auto flags = parse_list(+[](StringView x) { return String(x); });
            fetch_response.set_flags(move(flags));
            break;
        }
        case FetchCommand::DataItemType::InternalDate: {
            consume(" \"");
            auto date_view = parse_while([](u8 x) { return x != '"'; });
            consume("\"");
            auto date = Core::DateTime::parse("%d-%b-%Y %H:%M:%S %z", date_view).value();
            fetch_response.set_internal_date(date);
            break;
        }
        case FetchCommand::DataItemType::UID: {
            consume(" ");
            fetch_response.set_uid(parse_number());
            break;
        }
        case FetchCommand::DataItemType::PeekBody:
            // Spec doesn't allow for this in a response.
            m_parsing_failed = true;
            break;
        case FetchCommand::DataItemType::BodySection: {
            auto body = parse_nstring();
            fetch_response.add_body_data(move(data_item), body.has_value() ? body.release_value() : Optional<String>());
            break;
        }
        }
        if (!at_end() && m_buffer[position] != ')')
            consume(" ");
    }
    consume("\r\n");
    return fetch_response;
}

StringView Parser::parse_literal_string()
{
    consume("{");
    auto num_bytes = parse_number();
    consume("}\r\n");

    if (m_buffer.size() < position + num_bytes) {
        m_parsing_failed = true;
        return "";
    }

    position += num_bytes;
    return StringView(m_buffer.data() + position - num_bytes, num_bytes);
}

ListItem Parser::parse_list_item()
{
    consume(" ");
    auto flags_vec = parse_list(parse_mailbox_flag);
    unsigned flags = 0;
    for (auto flag : flags_vec) {
        flags |= static_cast<unsigned>(flag);
    }
    consume(" \"");
    auto reference = parse_while([](u8 x) { return x != '"'; });
    consume("\" ");
    auto mailbox = parse_astring();
    consume("\r\n");
    return ListItem { flags, String(reference), String(mailbox) };
}

void Parser::parse_capability_response()
{
    auto capability = AK::Vector<String>();
    while (!try_consume("\r\n")) {
        consume(" ");
        auto x = String(parse_atom());
        capability.append(x);
    }
    m_response.data().add_capabilities(move(capability));
}

StringView Parser::parse_atom()
{
    auto is_non_atom_char = [](u8 x) {
        auto non_atom_chars = { '(', ')', '{', ' ', '%', '*', '"', '\\', ']' };
        return AK::find(non_atom_chars.begin(), non_atom_chars.end(), x) != non_atom_chars.end();
    };

    auto start = position;
    auto count = 0;
    while (!at_end() && !is_ascii_control(m_buffer[position]) && !is_non_atom_char(m_buffer[position])) {
        count++;
        position++;
    }

    return StringView(m_buffer.data() + start, count);
}

ResponseStatus Parser::parse_status()
{
    auto atom = parse_atom();

    if (atom.matches("OK")) {
        return ResponseStatus::OK;
    } else if (atom.matches("BAD")) {
        return ResponseStatus::Bad;
    } else if (atom.matches("NO")) {
        return ResponseStatus::No;
    }

    m_parsing_failed = true;
    return ResponseStatus::Bad;
}

template<typename T>
Vector<T> Parser::parse_list(T converter(StringView))
{
    consume("(");
    Vector<T> x;
    bool first = true;
    while (!try_consume(")")) {
        if (!first)
            consume(" ");
        auto item = parse_while([](u8 x) {
            return x != ' ' && x != ')';
        });
        x.append(converter(item));
        first = false;
    }

    return x;
}

MailboxFlag Parser::parse_mailbox_flag(StringView s)
{
    if (s.matches("\\All"))
        return MailboxFlag::All;
    if (s.matches("\\Drafts"))
        return MailboxFlag::Drafts;
    if (s.matches("\\Flagged"))
        return MailboxFlag::Flagged;
    if (s.matches("\\HasChildren"))
        return MailboxFlag::HasChildren;
    if (s.matches("\\HasNoChildren"))
        return MailboxFlag::HasNoChildren;
    if (s.matches("\\Important"))
        return MailboxFlag::Important;
    if (s.matches("\\Junk"))
        return MailboxFlag::Junk;
    if (s.matches("\\Marked"))
        return MailboxFlag::Marked;
    if (s.matches("\\Noinferiors"))
        return MailboxFlag::NoInferiors;
    if (s.matches("\\Noselect"))
        return MailboxFlag::NoSelect;
    if (s.matches("\\Sent"))
        return MailboxFlag::Sent;
    if (s.matches("\\Trash"))
        return MailboxFlag::Trash;
    if (s.matches("\\Unmarked"))
        return MailboxFlag::Unmarked;

    dbgln("Unrecognized mailbox flag {}", s);
    return MailboxFlag::Unknown;
}

StringView Parser::parse_while(Function<bool(u8)> should_consume)
{
    int chars = 0;
    while (!at_end() && should_consume(m_buffer[position])) {
        position++;
        chars++;
    }
    return StringView(m_buffer.data() + position - chars, chars);
}

FetchCommand::DataItem Parser::parse_fetch_data_item()
{
    auto msg_attr = parse_while([](u8 x) { return is_ascii_alpha(x) != 0; });
    if (msg_attr.equals_ignoring_case("BODY") && try_consume("[")) {
        auto data_item = FetchCommand::DataItem {
            .type = FetchCommand::DataItemType::BodySection,
            .section = { {} }
        };
        auto section_type = parse_while([](u8 x) { return x != ']' && x != ' '; });
        if (section_type.equals_ignoring_case("HEADER.FIELDS")) {
            data_item.section->type = FetchCommand::DataItem::SectionType::HeaderFields;
            data_item.section->headers = Vector<String>();
            consume(" ");
            auto headers = parse_list(+[](StringView x) { return x; });
            for (auto& header : headers) {
                data_item.section->headers->append(header);
            }
            consume("]");
        } else if (section_type.equals_ignoring_case("HEADER.FIELDS.NOT")) {
            data_item.section->type = FetchCommand::DataItem::SectionType::HeaderFieldsNot;
            data_item.section->headers = Vector<String>();
            consume(" (");
            auto headers = parse_list(+[](StringView x) { return x; });
            for (auto& header : headers) {
                data_item.section->headers->append(header);
            }
            consume("]");
        } else if (is_ascii_digit(section_type[0])) {
            data_item.section->type = FetchCommand::DataItem::SectionType::Parts;
            data_item.section->parts = Vector<int>();

            while (!try_consume("]")) {
                auto num = parse_number();
                if (num != (unsigned)-1) {
                    data_item.section->parts->append((int)num);
                    continue;
                }
                auto atom = parse_atom();
                if (atom.equals_ignoring_case("MIME")) {
                    data_item.section->ends_with_mime = true;
                    continue;
                }
            }
        } else if (section_type.equals_ignoring_case("TEXT")) {
            data_item.section->type = FetchCommand::DataItem::SectionType::Text;
        } else if (section_type.equals_ignoring_case("HEADER")) {
            data_item.section->type = FetchCommand::DataItem::SectionType::Header;
        } else {
            dbgln("Unmatched section type {}", section_type);
            m_parsing_failed = true;
        }
        if (try_consume("<")) {
            auto start = parse_number();
            data_item.partial_fetch = true;
            data_item.start = (int)start;
            consume(">");
        }
        try_consume(" ");
        return data_item;
    } else if (msg_attr.equals_ignoring_case("FLAGS")) {
        return FetchCommand::DataItem {
            .type = FetchCommand::DataItemType::Flags
        };
    } else if (msg_attr.equals_ignoring_case("UID")) {
        return FetchCommand::DataItem {
            .type = FetchCommand::DataItemType::UID
        };
    } else if (msg_attr.equals_ignoring_case("INTERNALDATE")) {
        return FetchCommand::DataItem {
            .type = FetchCommand::DataItemType::InternalDate
        };
    } else if (msg_attr.equals_ignoring_case("ENVELOPE")) {
        return FetchCommand::DataItem {
            .type = FetchCommand::DataItemType::Envelope
        };
    } else {
        dbgln("msg_attr not matched: {}", msg_attr);
        m_parsing_failed = true;
        return FetchCommand::DataItem {};
    }
}
Optional<Vector<Address>> Parser::parse_address_list()
{
    if (try_consume("NIL"))
        return {};

    auto addresses = Vector<Address>();
    consume("(");
    while (!try_consume(")")) {
        addresses.append(parse_address());
        if (!at_end() && m_buffer[position] != ')')
            consume(" ");
    }
    return { addresses };
}

Address Parser::parse_address()
{
    consume("(");
    auto address = Address();
    // I hate this so much. Why is there no Optional.map??
    auto name = parse_nstring();
    address.name = name.has_value() ? Optional<String>(name.value()) : Optional<String>();
    consume(" ");
    auto source_route = parse_nstring();
    address.source_route = source_route.has_value() ? Optional<String>(source_route.value()) : Optional<String>();
    consume(" ");
    auto mailbox = parse_nstring();
    address.mailbox = mailbox.has_value() ? Optional<String>(mailbox.value()) : Optional<String>();
    consume(" ");
    auto host = parse_nstring();
    address.host = host.has_value() ? Optional<String>(host.value()) : Optional<String>();
    consume(")");
    return address;
}
StringView Parser::parse_astring()
{
    if (!at_end() && (m_buffer[position] == '{' || m_buffer[position] == '"'))
        return parse_string();
    else
        return parse_atom();
}
}