reactive-mail-server/backend/src/main/java/com/mailserver/imap/MimeTree.java

package com.mailserver.imap;

import jakarta.mail.MessagingException;
import jakarta.mail.internet.ContentType;
import jakarta.mail.internet.InternetHeaders;
import jakarta.mail.internet.ParseException;

import java.io.ByteArrayInputStream;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Enumeration;
import java.util.List;
import java.util.Locale;
import java.util.Set;

/**
 * A message split into MIME parts by byte offsets — what IMAP needs for BODYSTRUCTURE and BODY[section]:
 * every part is a range of the original bytes, so sizes and contents match BODY[] exactly.
 */
final class MimeTree {

    private static final int MAX_DEPTH = 20;

    final byte[] raw;
    /** Part = [start, end); headers = [start, bodyStart), the blank line included. */
    final int start;
    final int bodyStart;
    final int end;
    final InternetHeaders headers;
    final ContentType type;
    final List<MimeTree> children;
    /** The encapsulated message of a message/rfc822 part. */
    final MimeTree message;

    private MimeTree(byte[] raw, int start, int end, int depth) {
        this.raw = raw;
        this.start = start;
        this.end = end;
        this.bodyStart = headerEnd(raw, start, end);
        this.headers = parseHeaders(raw, start, bodyStart);
        this.type = contentType(headers);
        List<MimeTree> parts = new ArrayList<>();
        MimeTree nested = null;
        if (depth < MAX_DEPTH) {
            if (isMultipart()) {
                String boundary = type.getParameter("boundary");
                if (boundary != null && !boundary.isEmpty()) {
                    split(boundary, depth, parts);
                }
            } else if (is("message", "rfc822") || is("message", "global")) {
                nested = new MimeTree(raw, bodyStart, end, depth + 1);
            }
        }
        this.children = Collections.unmodifiableList(parts);
        this.message = nested;
    }

    static MimeTree parse(byte[] raw) {
        return new MimeTree(raw, 0, raw.length, 0);
    }

    boolean isMultipart() {
        return "multipart".equalsIgnoreCase(type.getPrimaryType());
    }

    boolean is(String primary, String sub) {
        return primary.equalsIgnoreCase(type.getPrimaryType()) && sub.equalsIgnoreCase(type.getSubType());
    }

    byte[] all() {
        return slice(start, end);
    }

    byte[] header() {
        return slice(start, bodyStart);
    }

    byte[] body() {
        return slice(bodyStart, end);
    }

    int bodySize() {
        return end - bodyStart;
    }

    /** Lines of the body (for text parts in BODYSTRUCTURE). */
    int bodyLines() {
        int lines = 0;
        for (int i = bodyStart; i < end; i++) {
            if (raw[i] == '\n') {
                lines++;
            }
        }
        if (end > bodyStart && raw[end - 1] != '\n') {
            lines++;
        }
        return lines;
    }

    /** Header value unfolded, or null. */
    String header(String name) {
        String[] values = headers.getHeader(name);
        if (values == null || values.length == 0 || values[0] == null) {
            return null;
        }
        return values[0].replaceAll("\r?\n(?=[ \t])", "").trim();
    }

    /** Only these header fields (or all but these), each with its folded lines, plus the blank line. */
    byte[] headerFields(Set<String> names, boolean not) {
        StringBuilder out = new StringBuilder();
        String text = new String(raw, start, bodyStart - start, StandardCharsets.ISO_8859_1);
        String[] lines = text.split("\r?\n", -1);
        boolean keep = false;
        for (String line : lines) {
            if (line.isEmpty()) {
                continue;
            }
            if (line.charAt(0) == ' ' || line.charAt(0) == '\t') {
                if (keep) {
                    out.append(line).append("\r\n");
                }
                continue;
            }
            int colon = line.indexOf(':');
            String name = colon > 0 ? line.substring(0, colon).trim().toUpperCase(Locale.ROOT) : "";
            keep = names.contains(name) != not;
            if (keep) {
                out.append(line).append("\r\n");
            }
        }
        out.append("\r\n");
        return out.toString().getBytes(StandardCharsets.ISO_8859_1);
    }

    private byte[] slice(int from, int to) {
        byte[] copy = new byte[Math.max(0, to - from)];
        System.arraycopy(raw, from, copy, 0, copy.length);
        return copy;
    }

    /** Body parts between "--boundary" lines; the CRLF before a delimiter belongs to the delimiter. */
    private void split(String boundary, int depth, List<MimeTree> parts) {
        byte[] delimiter = ("--" + boundary).getBytes(StandardCharsets.ISO_8859_1);
        int partStart = -1;
        int i = bodyStart;
        while (i < end) {
            int lineEnd = indexOf(raw, (byte) '\n', i, end);
            int next = lineEnd < 0 ? end : lineEnd + 1;
            if (startsWith(raw, i, end, delimiter)) {
                int after = i + delimiter.length;
                boolean close = after + 1 < end && raw[after] == '-' && raw[after + 1] == '-';
                if (partStart >= 0) {
                    int partEnd = i;
                    if (partEnd > partStart && raw[partEnd - 1] == '\n') {
                        partEnd--;
                        if (partEnd > partStart && raw[partEnd - 1] == '\r') {
                            partEnd--;
                        }
                    }
                    parts.add(new MimeTree(raw, partStart, Math.max(partStart, partEnd), depth + 1));
                }
                if (close) {
                    return;
                }
                partStart = next;
            }
            i = next;
        }
        if (partStart >= 0 && partStart < end) { // no closing delimiter — take the rest
            parts.add(new MimeTree(raw, partStart, end, depth + 1));
        }
    }

    private static int headerEnd(byte[] raw, int start, int end) {
        if (start < end && raw[start] == '\n') {
            return start + 1;
        }
        if (start + 1 < end && raw[start] == '\r' && raw[start + 1] == '\n') {
            return start + 2;
        }
        for (int i = start; i < end; i++) {
            if (raw[i] == '\n') {
                if (i + 1 < end && raw[i + 1] == '\n') {
                    return i + 2;
                }
                if (i + 2 < end && raw[i + 1] == '\r' && raw[i + 2] == '\n') {
                    return i + 3;
                }
            }
        }
        return end; // headers only
    }

    private static InternetHeaders parseHeaders(byte[] raw, int start, int end) {
        try {
            return new InternetHeaders(new ByteArrayInputStream(raw, start, end - start));
        } catch (MessagingException e) {
            return new InternetHeaders();
        }
    }

    private static ContentType contentType(InternetHeaders headers) {
        String[] values = headers.getHeader("Content-Type");
        if (values != null && values.length > 0 && values[0] != null) {
            try {
                return new ContentType(values[0].replaceAll("\r?\n(?=[ \t])", "").trim());
            } catch (ParseException ignored) {
                // broken header — the RFC 2045 default
            }
        }
        try {
            return new ContentType("text/plain; charset=us-ascii");
        } catch (ParseException e) {
            throw new IllegalStateException(e);
        }
    }

    private static boolean startsWith(byte[] raw, int at, int end, byte[] prefix) {
        if (at + prefix.length > end) {
            return false;
        }
        for (int i = 0; i < prefix.length; i++) {
            if (raw[at + i] != prefix[i]) {
                return false;
            }
        }
        return true;
    }

    private static int indexOf(byte[] raw, byte value, int from, int end) {
        for (int i = from; i < end; i++) {
            if (raw[i] == value) {
                return i;
            }
        }
        return -1;
    }

    /** Header names of the part, for tests and debugging. */
    List<String> headerNames() {
        List<String> names = new ArrayList<>();
        Enumeration<jakarta.mail.Header> all = headers.getAllHeaders();
        while (all.hasMoreElements()) {
            names.add(all.nextElement().getName());
        }
        return names;
    }
}