From 8887b9b95138f70b8c26f6278d8bd1551c45048c Mon Sep 17 00:00:00 2001 From: Ivars Belovs Date: Mon, 24 Aug 2026 14:18:23 +0300 Subject: [PATCH 1/3] Fetch message parts individually instead of buffering RFC822 get_message fetched RFC822, which materializes the entire message as one String, then parsed it with Mail and base64-decoded each attachment while the raw message was still reachable. Measured on an 8.27MB message, peak live strings reached 55.6MB (6.7x the wire size): net-imap literals arrive as ASCII-8BIT, and Mail.new on a binary source eagerly allocates two more full-size copies on top of the aliased raw_source, then materializing the part tree copies every part's base64 body again. Worse, RSS never came back. Twelve sequential fetches of a 25MB-attachment message grew RSS by one wire size each time (+34.5MB/fetch, 175MB to 543MB) with the Ruby heap provably flat -- 47,579 live objects and 0.42MB of strings, unchanged. Plain large alloc/free plateaus, so this is fragmentation specific to the Mail parse path rather than generic allocator behaviour. In production that ratchet OOMKilled the pod every ~10 minutes at a 384Mi limit. Now MessageReader walks BODYSTRUCTURE, fetches the header and inline bodies on their own, and streams each attachment with partial fetches (BODY.PEEK[n]) through an incremental base64 decoder into a Tempfile, uploading and discarding it before the next attachment is fetched. Peak memory tracks the 4MiB chunk rather than the message: RSS is flat across 12 fetches, and five concurrent 25MB-attachment fetches peak at 276MB against 614MB before. Verified with a differential harness over 11 fixtures (single-part, alternative, nested related>alternative, non-UTF-8 charset, quoted-printable, attachment without a filename, inline image with a Content-ID, message/rfc822, multi-chunk) comparing against the old whole-message output. Attachment bytes are identical in every case. Three intentional differences: - text_body on a non-multipart message. Mail#text_part searches all_parts, which is empty unless the message is multipart, so the old code returned nil for the body of every plain-text email. It is now returned. - content_type is the bare MIME type rather than the whole Content-Type header; filename is already its own field. - a 7bit/8bit attachment keeps the line endings the server sent, where Mail#decoded normalized CRLF to LF. Body sections are requested with BODY.PEEK so reading cannot set \Seen, and inline parts are rebuilt from a synthesized MIME header so Mail still applies the charset conversion. AttachmentStore gains a streaming upload_io and memoizes its client and presigner, which were being rebuilt twice per attachment at ~11MB of RSS each. Deployment note: attachments are now spooled through Tempfile, so TMPDIR must point at disk. If the container's /tmp is tmpfs the spool is anonymous memory and counts against the cgroup limit, which would defeat the fix. Co-Authored-By: Claude Opus 5 (1M context) --- lib/mail_mcp.rb | 3 + lib/mail_mcp/attachment_store.rb | 64 ++++++---- lib/mail_mcp/base64_stream.rb | 31 +++++ lib/mail_mcp/imap_client.rb | 47 +------- lib/mail_mcp/message_reader.rb | 151 ++++++++++++++++++++++++ lib/mail_mcp/message_structure.rb | 75 ++++++++++++ spec/mail_mcp/attachment_store_spec.rb | 25 ++++ spec/mail_mcp/base64_stream_spec.rb | 51 ++++++++ spec/mail_mcp/imap_client_spec.rb | 120 +++++++++++++++---- spec/mail_mcp/message_structure_spec.rb | 114 ++++++++++++++++++ spec/spec_helper.rb | 2 + spec/support/fake_imap_server.rb | 122 +++++++++++++++++++ 12 files changed, 721 insertions(+), 84 deletions(-) create mode 100644 lib/mail_mcp/base64_stream.rb create mode 100644 lib/mail_mcp/message_reader.rb create mode 100644 lib/mail_mcp/message_structure.rb create mode 100644 spec/mail_mcp/base64_stream_spec.rb create mode 100644 spec/mail_mcp/message_structure_spec.rb create mode 100644 spec/support/fake_imap_server.rb diff --git a/lib/mail_mcp.rb b/lib/mail_mcp.rb index 7794877..03892c7 100644 --- a/lib/mail_mcp.rb +++ b/lib/mail_mcp.rb @@ -16,6 +16,9 @@ class << self require_relative "mail_mcp/jwt_service" require_relative "mail_mcp/pkce" +require_relative "mail_mcp/base64_stream" +require_relative "mail_mcp/message_structure" +require_relative "mail_mcp/message_reader" require_relative "mail_mcp/imap_client" require_relative "mail_mcp/smtp_client" require_relative "mail_mcp/attachment_store" diff --git a/lib/mail_mcp/attachment_store.rb b/lib/mail_mcp/attachment_store.rb index bd962a0..24c96d2 100644 --- a/lib/mail_mcp/attachment_store.rb +++ b/lib/mail_mcp/attachment_store.rb @@ -5,28 +5,52 @@ module MailMCP module AttachmentStore EXPIRY = 7 * 24 * 3600 - def self.upload(content:, filename:, content_type:) - key = "attachments/#{SecureRandom.uuid}/#{filename}" - bucket = ENV.fetch("AWS_S3_BUCKET") - - s3.put_object( - bucket: bucket, - key: key, - body: content, - content_type: content_type - ) - - presigner.presigned_url(:get_object, bucket: bucket, key: key, expires_in: EXPIRY) - end + class << self + # Uploads an in-memory body. Prefer #upload_io for anything attachment-sized: + # a String body has to be fully resident, which is what we are trying to avoid. + def upload(content:, filename:, content_type:) + store(body: content, filename: filename, content_type: content_type) + end - def self.s3 - Aws::S3::Client.new - end - private_class_method :s3 + # Uploads from an open IO (typically a Tempfile), letting the SDK read it in + # chunks instead of holding the whole attachment as a String. + def upload_io(io:, filename:, content_type:) + io.rewind + store(body: io, filename: filename, content_type: content_type) + end + + # Drops the memoized clients. Only needed by tests and after credential changes. + def reset! + @s3 = nil + @presigner = nil + end + + private + + def store(body:, filename:, content_type:) + key = "attachments/#{SecureRandom.uuid}/#{filename}" + bucket = ENV.fetch("AWS_S3_BUCKET") + + s3.put_object( + bucket: bucket, + key: key, + body: body, + content_type: content_type + ) + + presigner.presigned_url(:get_object, bucket: bucket, key: key, expires_in: EXPIRY) + end + + # Memoized: building a client cost ~11MB of RSS, and the un-memoized version + # built two per attachment — once here and once more inside #presigner. + # Aws::S3::Client is thread-safe, so one instance serves every Puma thread. + def s3 + @s3 ||= Aws::S3::Client.new + end - def self.presigner - Aws::S3::Presigner.new(client: s3) + def presigner + @presigner ||= Aws::S3::Presigner.new(client: s3) + end end - private_class_method :presigner end end diff --git a/lib/mail_mcp/base64_stream.rb b/lib/mail_mcp/base64_stream.rb new file mode 100644 index 0000000..b565d80 --- /dev/null +++ b/lib/mail_mcp/base64_stream.rb @@ -0,0 +1,31 @@ +module MailMCP + # Incremental base64 decoder. Feed it arbitrary slices of a base64 stream and it + # emits decoded bytes as soon as it holds complete four-character groups, carrying + # the remainder (always fewer than four characters) into the next call. This is what + # lets an attachment of any size be decoded with a bounded amount of memory. + class Base64Stream + GROUP = 4 + + def initialize + @carry = +"" + end + + def push(chunk) + buffer = @carry + chunk.delete("\r\n\t ") + complete = buffer.bytesize - (buffer.bytesize % GROUP) + @carry = buffer.byteslice(complete, buffer.bytesize - complete) || +"" + return "".b if complete.zero? + + buffer.byteslice(0, complete).unpack1("m") + end + + # Decodes whatever is left over, including any padding. + def finish + return "".b if @carry.empty? + + remainder = @carry + @carry = +"" + remainder.unpack1("m") + end + end +end diff --git a/lib/mail_mcp/imap_client.rb b/lib/mail_mcp/imap_client.rb index 9b12886..fd19b52 100644 --- a/lib/mail_mcp/imap_client.rb +++ b/lib/mail_mcp/imap_client.rb @@ -75,16 +75,14 @@ def list_messages(folder:, page: 1, per_page: 20) { messages: messages, total: total, page: page, per_page: per_page } end + # Delegates to MessageReader, which fetches the message a part at a time so peak + # memory tracks the largest part rather than the whole message. def get_message(folder:, uid:) MailMCP.logger.info { "IMAP get_message folder=#{folder.inspect} uid=#{uid}" } @imap.examine(folder) - data = @imap.uid_fetch([uid.to_i], %w[RFC822 FLAGS]).first - unless data - MailMCP.logger.warn { "IMAP get_message not found folder=#{folder.inspect} uid=#{uid}" } - return nil - end - - format_message(uid: uid, parsed: Mail.new(data.attr["RFC822"]), flags: data.attr["FLAGS"]) + message = MessageReader.new(@imap, uid.to_i).read + MailMCP.logger.warn { "IMAP get_message not found folder=#{folder.inspect} uid=#{uid}" } unless message + message end def search_messages(folder:, query:) @@ -142,27 +140,6 @@ def self.open_connection(config) private - def format_message(uid:, parsed:, flags:) - { - uid: uid, - message_id: parsed.message_id, - in_reply_to: parsed.in_reply_to, - references: parsed.references, - subject: parsed.subject, - from: parsed.from, - sender: parsed.sender, - reply_to: parsed.reply_to, - to: parsed.to, - cc: parsed.cc, - bcc: parsed.bcc, - date: parsed.date&.iso8601, - text_body: parsed.text_part&.decoded, - html_body: parsed.html_part&.decoded, - flags: flags, - attachments: extract_attachments(parsed) - } - end - def format_envelope(msg) env = msg.attr["ENVELOPE"] { @@ -187,19 +164,5 @@ def format_addresses(addrs) addrs.map { |a| "#{a.name} <#{a.mailbox}@#{a.host}>" } end - - def extract_attachments(mail) - mail.attachments.map do |att| - # Mail#decoded does not memoize — it base64-decodes into a new String on - # every call, so decode once and reuse for both the upload and the size. - content = att.decoded - url = AttachmentStore.upload( - content: content, - filename: att.filename || "attachment", - content_type: att.content_type - ) - { filename: att.filename, content_type: att.content_type, size: content.bytesize, url: url } - end - end end end diff --git a/lib/mail_mcp/message_reader.rb b/lib/mail_mcp/message_reader.rb new file mode 100644 index 0000000..8e86841 --- /dev/null +++ b/lib/mail_mcp/message_reader.rb @@ -0,0 +1,151 @@ +require "tempfile" + +module MailMCP + # Reads one message over an existing IMAP connection a part at a time, so peak memory + # tracks the part being handled rather than the size of the whole message. Fetching + # RFC822 instead would materialize the entire message as a single String. + class MessageReader + # Encoded bytes pulled per FETCH when streaming an attachment. Bounds peak memory + # regardless of attachment size; larger chunks mean fewer round trips but hold a + # Puma thread for longer. + CHUNK_SIZE = 4 * 1024 * 1024 + + def initialize(imap, uid) + @imap = imap + @uid = uid + end + + # Returns the message as the tool-facing hash, or nil when the uid is gone. + def read + data = @imap.uid_fetch([@uid], ["BODYSTRUCTURE", "FLAGS", "BODY.PEEK[HEADER]"])&.first + return nil unless data + + parts = MessageStructure.flatten(data.attr["BODYSTRUCTURE"]) + MailMCP.logger.debug { "IMAP message uid=#{@uid} parts=#{parts.map(&:section).inspect}" } + + format( + headers: Mail.new(self.class.body_of(data.attr, "HEADER").to_s), + flags: data.attr["FLAGS"], + parts: parts + ) + end + + # Servers answer BODY.PEEK[...] with BODY[...], and append the origin octet on a + # partial fetch (BODY[1]<0>), so the response key never matches the request. + def self.body_of(attrs, section) + exact = attrs["BODY[#{section}]"] + return exact if exact + + prefix = "BODY[#{section}]<" + attrs.find { |key, _value| key.to_s.start_with?(prefix) }&.last + end + + private + + def format(headers:, flags:, parts:) + inline = parts.reject(&:attachment?) + { + uid: @uid, + message_id: headers.message_id, + in_reply_to: headers.in_reply_to, + references: headers.references, + subject: headers.subject, + from: headers.from, + sender: headers.sender, + reply_to: headers.reply_to, + to: headers.to, + cc: headers.cc, + bcc: headers.bcc, + date: headers.date&.iso8601, + text_body: inline_body(inline, "text/plain"), + html_body: inline_body(inline, "text/html"), + flags: flags, + attachments: upload_attachments(parts.select(&:attachment?)) + } + end + + # Mail#text_part/#html_part pick the first non-attachment part of that type, so + # match the flattened list the same way. + def inline_body(inline, mime_type) + part = inline.find { |candidate| candidate.mime_type == mime_type } + return nil unless part + + body = fetch_section(part.section) + body && decode_inline(part, body) + end + + # Rebuilding the part from a synthesized MIME header lets Mail apply both the + # transfer-encoding and the charset conversion, which is what #decoded did on the + # old whole-message parse. Hand-rolling that would mislabel any non-UTF-8 body. + def decode_inline(part, body) + header = "Content-Type: #{part.mime_type}" + header += "; charset=#{part.charset}" if part.charset + header += "\r\nContent-Transfer-Encoding: #{part.encoding}\r\n" + Mail::Part.new("#{header}\r\n#{body}").decoded + rescue StandardError => e + MailMCP.logger.warn { "IMAP inline decode failed section=#{part.section}: #{e.class}: #{e.message}" } + nil + end + + # Each attachment is streamed to disk, uploaded, then discarded before the next is + # fetched, so attachments never accumulate in memory. + def upload_attachments(parts) + parts.map do |part| + tempfile = fetch_attachment(part) + begin + url = AttachmentStore.upload_io( + io: tempfile, + filename: part.filename || "attachment", + content_type: part.mime_type + ) + { filename: part.filename, content_type: part.mime_type, size: tempfile.size, url: url } + ensure + tempfile.close! + end + end + end + + def fetch_attachment(part) + tempfile = Tempfile.new("mail_mcp_attachment") + tempfile.binmode + if part.streamable? + stream_section(tempfile, part) + else + # Quoted-printable, or a part whose size the server did not report. Neither is + # used for large payloads in practice, so one pass is acceptable here. + tempfile.write(decode_whole(fetch_section(part.section).to_s, part.encoding)) + end + tempfile.flush + tempfile + end + + def stream_section(tempfile, part) + decoder = part.base64? ? Base64Stream.new : nil + offset = 0 + # Bounded by the size BODYSTRUCTURE reported; the empty check is a second guard so + # a server returning short reads cannot spin here. + while offset < part.encoded_size + chunk = fetch_section(part.section, offset: offset, length: CHUNK_SIZE) + break if chunk.nil? || chunk.empty? + + tempfile.write(decoder ? decoder.push(chunk) : chunk) + offset += chunk.bytesize + end + tempfile.write(decoder.finish) if decoder + end + + def decode_whole(raw, encoding) + encoder = Mail::Encodings.get_encoding(encoding) + encoder ? encoder.decode(raw) : raw + rescue StandardError => e + MailMCP.logger.warn { "IMAP attachment decode failed encoding=#{encoding.inspect}: #{e.class}: #{e.message}" } + raw + end + + def fetch_section(section, offset: nil, length: nil) + spec = offset ? "BODY.PEEK[#{section}]<#{offset}.#{length}>" : "BODY.PEEK[#{section}]" + data = @imap.uid_fetch([@uid], [spec])&.first + data && self.class.body_of(data.attr, section) + end + end +end diff --git a/lib/mail_mcp/message_structure.rb b/lib/mail_mcp/message_structure.rb new file mode 100644 index 0000000..aaedac7 --- /dev/null +++ b/lib/mail_mcp/message_structure.rb @@ -0,0 +1,75 @@ +require "net/imap" + +module MailMCP + # Flattens an IMAP BODYSTRUCTURE into its leaf parts, each tagged with the + # RFC 3501 section number needed to FETCH that part on its own. Fetching parts + # individually is what lets a large message be handled without ever holding the + # whole thing in memory. + module MessageStructure + IDENTITY_ENCODINGS = ["7bit", "8bit", "binary", ""].freeze + + Part = Struct.new(:section, :media_type, :subtype, :encoding, :encoded_size, :filename, + :charset, :content_id, keyword_init: true) do + def mime_type + "#{media_type}/#{subtype}" + end + + # Mirrors Mail::Message#attachment?, which keys off a filename being present + # rather than off the Content-Disposition type. + def attachment? + !filename.nil? + end + + def base64? + encoding == "base64" + end + + # Encodings we can fetch and write through in chunks. Anything else (in + # practice quoted-printable) has to be decoded in one pass. + def streamable? + encoded_size.positive? && (base64? || IDENTITY_ENCODINGS.include?(encoding)) + end + end + + class << self + def flatten(body, prefix = nil) + return [] if body.nil? + + if body.multipart? + body.parts.each_with_index.flat_map do |child, index| + flatten(child, [prefix, index + 1].compact.join(".")) + end + else + [leaf(body, prefix || "1")] + end + end + + private + + def leaf(body, section) + Part.new( + section: section, + media_type: body.media_type.to_s.downcase, + subtype: body.subtype.to_s.downcase, + encoding: body.encoding.to_s.downcase, + encoded_size: body.size.to_i, + filename: filename_for(body), + charset: param(body.param, "charset"), + content_id: body.content_id + ) + end + + def filename_for(body) + param(body.disposition&.param, "filename") || param(body.param, "name") + end + + # Servers pick their own case for parameter names, so never index directly. + def param(params, name) + return nil unless params + + pair = params.find { |key, _value| key.to_s.casecmp?(name) } + pair&.last + end + end + end +end diff --git a/spec/mail_mcp/attachment_store_spec.rb b/spec/mail_mcp/attachment_store_spec.rb index 06ad18a..880f928 100644 --- a/spec/mail_mcp/attachment_store_spec.rb +++ b/spec/mail_mcp/attachment_store_spec.rb @@ -11,6 +11,8 @@ allow(s3_client).to receive(:put_object) allow(presigner).to receive(:presigned_url).and_return(presigned_url) stub_const("ENV", ENV.to_h.merge("AWS_S3_BUCKET" => "test-bucket")) + # The client and presigner are memoized to avoid rebuilding them per attachment. + described_class.reset! end describe ".upload" do @@ -26,6 +28,12 @@ expect(url).to eq(presigned_url) end + it "reuses one memoized client and presigner across uploads" do + 2.times { described_class.upload(content: "d", filename: "f.txt", content_type: "text/plain") } + expect(Aws::S3::Client).to have_received(:new).once + expect(Aws::S3::Presigner).to have_received(:new).once + end + it "generates a unique S3 key for each upload" do keys = 2.times.map do key = nil @@ -36,4 +44,21 @@ expect(keys.first).not_to eq(keys.last) end end + + describe ".upload_io" do + it "streams the IO to S3 as the request body" do + io = StringIO.new("streamed bytes") + url = described_class.upload_io(io: io, filename: "big.bin", content_type: "application/octet-stream") + + expect(s3_client).to have_received(:put_object).with(hash_including(body: io)) + expect(url).to eq(presigned_url) + end + + it "rewinds the IO so a partially read tempfile still uploads in full" do + io = StringIO.new("streamed bytes") + io.read(4) + described_class.upload_io(io: io, filename: "big.bin", content_type: "application/octet-stream") + expect(io.pos).to eq(0) + end + end end diff --git a/spec/mail_mcp/base64_stream_spec.rb b/spec/mail_mcp/base64_stream_spec.rb new file mode 100644 index 0000000..df5bd8c --- /dev/null +++ b/spec/mail_mcp/base64_stream_spec.rb @@ -0,0 +1,51 @@ +require "spec_helper" + +RSpec.describe MailMCP::Base64Stream do + def decode_in_slices(payload, slice_size) + encoded = [payload].pack("m") + stream = described_class.new + out = +"" + encoded.scan(/.{1,#{slice_size}}/m) { |slice| out << stream.push(slice) } + out << stream.finish + out + end + + it "decodes a stream fed in one piece" do + expect(decode_in_slices("hello world", 4096)).to eq("hello world") + end + + # The point of the class: chunk boundaries almost never land on a 4-character + # base64 group, so the remainder has to survive between calls. + it "decodes identically no matter where the chunks split" do + payload = (0..255).to_a.pack("C*") * 40 + + [1, 2, 3, 5, 7, 16, 64, 1000].each do |slice_size| + expect(decode_in_slices(payload, slice_size)).to eq(payload), "failed at slice size #{slice_size}" + end + end + + it "handles payloads at each padding length" do + (0..3).each do |extra| + payload = ("x" * 30) + ("y" * extra) + expect(decode_in_slices(payload, 5)).to eq(payload) + end + end + + it "ignores the line breaks that wrap base64 in real messages" do + payload = "binary\x00\xFFdata".b * 50 + wrapped = [payload].pack("m") # pack("m") already wraps at 60 characters + stream = described_class.new + expect(stream.push(wrapped) + stream.finish).to eq(payload) + end + + it "returns empty strings rather than nil when it has nothing to emit yet" do + stream = described_class.new + expect(stream.push("ab")).to eq("") + expect(described_class.new.finish).to eq("") + end + + it "preserves binary encoding" do + payload = "\x00\x01\xFE\xFF".b * 100 + expect(decode_in_slices(payload, 9).b).to eq(payload) + end +end diff --git a/spec/mail_mcp/imap_client_spec.rb b/spec/mail_mcp/imap_client_spec.rb index d2b3985..7e737b5 100644 --- a/spec/mail_mcp/imap_client_spec.rb +++ b/spec/mail_mcp/imap_client_spec.rb @@ -52,6 +52,7 @@ describe "#get_message" do let(:attachment_body) { "PDF-CONTENT" * 10 } + let(:uploaded) { [] } let(:raw_message) do mail = Mail.new mail.from = "alice@example.com" @@ -61,49 +62,124 @@ mail.add_file(filename: "invoice.pdf", content: attachment_body) mail.to_s end + let(:server) { FakeImapServer.new(raw_message) } before do - fetch_data = instance_double( - Net::IMAP::FetchData, - attr: { "RFC822" => raw_message, "FLAGS" => [:Seen] } - ) - allow(imap).to receive(:uid_fetch).and_return([fetch_data]) - allow(MailMCP::AttachmentStore).to receive(:upload).and_return("https://s3.example.com/invoice.pdf") + allow(MailMCP::AttachmentStore).to receive(:upload_io) do |io:, filename:, **| + io.rewind + uploaded << { filename: filename, bytes: io.read } + "https://s3.example.com/#{filename}" + end end it "returns the message with its attachment metadata" do - result = described_class.new(imap).get_message(folder: "INBOX", uid: 42) + result = described_class.new(server).get_message(folder: "INBOX", uid: 42) expect(result[:subject]).to eq("Invoice") + expect(result[:text_body]).to eq("see attached") expect(result[:attachments]).to contain_exactly( { filename: "invoice.pdf", - content_type: a_string_including("application/pdf"), + content_type: "application/pdf", size: attachment_body.bytesize, url: "https://s3.example.com/invoice.pdf" } ) end - # Mail#decoded re-decodes on every call, so decoding twice doubled the peak - # memory of the largest allocation in the request path. - it "decodes each attachment only once" do - decode_count = 0 - allow_any_instance_of(Mail::Part).to receive(:decoded).and_wrap_original do |original| # rubocop:disable RSpec/AnyInstance - # Count only the attachment: format_message legitimately decodes - # text_part/html_part too, and those are Mail::Part instances as well. - decode_count += 1 if original.receiver.attachment? - original.call - end + it "uploads the attachment bytes intact" do + described_class.new(server).get_message(folder: "INBOX", uid: 42) + expect(uploaded.first[:bytes]).to eq(attachment_body) + end - described_class.new(imap).get_message(folder: "INBOX", uid: 42) + # The whole point of the per-part path: RFC822 pulls the entire message into one + # String, which is what made peak memory scale with message size. + it "never fetches the whole message" do + described_class.new(server).get_message(folder: "INBOX", uid: 42) + sections = server.fetches.map(&:first) + expect(sections).to include("HEADER") + expect(sections).not_to include("", "TEXT") + end + + # BODY[...] implicitly sets \Seen; BODY.PEEK[...] does not. + it "requests every body section with PEEK so reading cannot set flags" do + described_class.new(server).get_message(folder: "INBOX", uid: 42) + body_specs = server.requested_specs.grep(/\ABODY(\.PEEK)?\[/) - expect(decode_count).to eq(1) + expect(body_specs).not_to be_empty + expect(body_specs).to all(start_with("BODY.PEEK[")) end it "returns nil when the uid is not found" do - allow(imap).to receive(:uid_fetch).and_return([]) - expect(described_class.new(imap).get_message(folder: "INBOX", uid: 99)).to be_nil + allow(server).to receive(:uid_fetch).and_return([]) + expect(described_class.new(server).get_message(folder: "INBOX", uid: 99)).to be_nil + end + + context "with a single-part message" do + let(:raw_message) do + mail = Mail.new + mail.from = "alice@example.com" + mail.to = "bob@example.com" + mail.subject = "Plain" + mail.body = "just text" + mail.to_s + end + + # Mail#text_part searches all_parts, which is empty for a non-multipart message, + # so the previous whole-message implementation returned no body at all here. + it "returns the body" do + result = described_class.new(server).get_message(folder: "INBOX", uid: 42) + expect(result[:text_body]).to eq("just text") + expect(result[:attachments]).to be_empty + end + end + + context "with a non-UTF-8 body" do + let(:raw_message) do + mail = Mail.new + mail.from = "alice@example.com" + mail.to = "bob@example.com" + mail.subject = "Latin" + mail.content_type = "text/plain; charset=ISO-8859-1" + mail.content_transfer_encoding = "8bit" + mail.body = "caf\xE9 cr\xE8me".dup.force_encoding("ASCII-8BIT") + mail.to_s + end + + it "converts the body to UTF-8" do + result = described_class.new(server).get_message(folder: "INBOX", uid: 42) + expect(result[:text_body]).to eq("café crème") + end + end + + context "with an attachment larger than one chunk" do + let(:attachment_body) { "0123456789abcdef" * 700_000 } + + it "fetches it in byte ranges and reassembles it exactly" do + described_class.new(server).get_message(folder: "INBOX", uid: 42) + + expect(server.partial_fetches.size).to be > 1 + expect(uploaded.first[:bytes].bytesize).to eq(attachment_body.bytesize) + expect(uploaded.first[:bytes]).to eq(attachment_body) + end + end + + context "with a quoted-printable attachment" do + let(:raw_message) do + mail = Mail.new + mail.from = "alice@example.com" + mail.to = "bob@example.com" + mail.subject = "QP" + mail.text_part = Mail::Part.new { body "body" } + mail.attachments["notes.txt"] = { content: "line = one\r\n" * 20, + transfer_encoding: "quoted-printable" } + mail.to_s + end + + it "decodes it in a single pass" do + described_class.new(server).get_message(folder: "INBOX", uid: 42) + expect(uploaded.first[:bytes]).to include("line = one") + end end end diff --git a/spec/mail_mcp/message_structure_spec.rb b/spec/mail_mcp/message_structure_spec.rb new file mode 100644 index 0000000..a0d7ecc --- /dev/null +++ b/spec/mail_mcp/message_structure_spec.rb @@ -0,0 +1,114 @@ +require "spec_helper" + +RSpec.describe MailMCP::MessageStructure do + def structure_for(raw) + server = FakeImapServer.new(raw) + server.uid_fetch([1], ["BODYSTRUCTURE"]).first.attr["BODYSTRUCTURE"] + end + + def parts_for(raw) + described_class.flatten(structure_for(raw)) + end + + it "returns nothing for a nil structure" do + expect(described_class.flatten(nil)).to eq([]) + end + + # RFC 3501: a non-multipart message's body is section 1. + it "numbers a single-part message as section 1" do + mail = Mail.new + mail.from = "a@x.com" + mail.to = "b@x.com" + mail.body = "text" + raw = mail.to_s + expect(parts_for(raw).map(&:section)).to eq(["1"]) + end + + it "numbers the children of a top-level multipart without adding a level" do + raw = Mail.new do + from "a@x.com" + to "b@x.com" + text_part { body "plain" } + html_part { body "

html

" } + end.to_s + + parts = parts_for(raw) + expect(parts.map(&:section)).to eq(%w[1 2]) + expect(parts.map(&:mime_type)).to eq(["text/plain", "text/html"]) + end + + it "numbers nested multiparts with dotted sections" do + inner = Mail::Part.new do + content_type "multipart/alternative" + text_part { body "plain" } + html_part { body "

html

" } + end + outer = Mail.new do + from "a@x.com" + to "b@x.com" + content_type "multipart/mixed" + end + outer.add_part(inner) + outer.add_file(filename: "a.pdf", content: "PDF") + + expect(parts_for(outer.to_s).map(&:section)).to eq(%w[1.1 1.2 2]) + end + + describe "attachment detection" do + it "treats a part with a filename as an attachment" do + mail = Mail.new do + from "a@x.com" + to "b@x.com" + text_part { body "see attached" } + end + mail.add_file(filename: "invoice.pdf", content: "PDF") + + parts = parts_for(mail.to_s) + expect(parts.reject(&:attachment?).map(&:mime_type)).to eq(["text/plain"]) + expect(parts.select(&:attachment?).map(&:filename)).to eq(["invoice.pdf"]) + end + + it "reads a filename that the server reported in upper case" do + mail = Mail.new do + from "a@x.com" + to "b@x.com" + text_part { body "body" } + end + mail.add_file(filename: "report.csv", content: "a,b") + + # FakeImapServer upcases parameter names the way real servers do. + expect(parts_for(mail.to_s).find(&:attachment?).filename).to eq("report.csv") + end + end + + describe "streamability" do + it "streams base64 and identity encodings" do + %w[base64 7bit 8bit binary].each do |encoding| + part = described_class::Part.new(section: "1", encoding: encoding, encoded_size: 10) + expect(part).to be_streamable, "expected #{encoding} to stream" + end + end + + it "does not stream quoted-printable, which needs a single-pass decode" do + part = described_class::Part.new(section: "1", encoding: "quoted-printable", encoded_size: 10) + expect(part).not_to be_streamable + end + + # Without a size from the server the chunk loop has no bound to walk. + it "does not stream a part of unreported size" do + part = described_class::Part.new(section: "1", encoding: "base64", encoded_size: 0) + expect(part).not_to be_streamable + end + end + + it "exposes the charset so the body can be converted to UTF-8" do + raw = Mail.new do + from "a@x.com" + to "b@x.com" + content_type "text/plain; charset=ISO-8859-1" + body "text" + end.to_s + + expect(parts_for(raw).first.charset).to eq("ISO-8859-1") + end +end diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb index 0e61df4..d19d116 100644 --- a/spec/spec_helper.rb +++ b/spec/spec_helper.rb @@ -13,6 +13,8 @@ MailMCP.logger = Logger.new(IO::NULL) +require_relative "support/fake_imap_server" + RSpec.configure do |config| config.expect_with :rspec do |expectations| expectations.include_chain_clauses_in_custom_matcher_descriptions = true diff --git a/spec/support/fake_imap_server.rb b/spec/support/fake_imap_server.rb new file mode 100644 index 0000000..ce59c87 --- /dev/null +++ b/spec/support/fake_imap_server.rb @@ -0,0 +1,122 @@ +require "net/imap" + +# Serves BODYSTRUCTURE and BODY[...] fetches for a real RFC822 message, the way an +# IMAP server would. Addressing is derived from Mail's own part tree rather than from +# MailMCP::MessageStructure, so specs exercise the walker instead of confirming itself. +class FakeImapServer + attr_reader :fetches, :requested_specs + + def initialize(raw) + @mail = Mail.new(raw) + @fetches = [] + @requested_specs = [] + end + + def examine(_folder) = nil + + def select(_folder) = nil + + def uid_fetch(_uids, specs) + @requested_specs.concat(specs) + attrs = specs.each_with_object({}) { |spec, acc| serve(spec, acc) } + return nil if attrs.empty? + + [Struct.new(:attr).new(attrs)] + end + + # Sections fetched with a byte range, e.g. [["2", 0, 4194304]]. + def partial_fetches + @fetches.select { |_section, offset, _length| offset } + end + + private + + def serve(spec, acc) + case spec + when "BODYSTRUCTURE" then acc["BODYSTRUCTURE"] = structure_of(@mail) + when "FLAGS" then acc["FLAGS"] = [:Seen] + when /\ABODY\.PEEK\[(.+?)\](?:<(\d+)\.(\d+)>)?\z/ + serve_body(acc, Regexp.last_match(1), Regexp.last_match(2)&.to_i, Regexp.last_match(3)&.to_i) + end + end + + def serve_body(acc, section, offset, length) + @fetches << [section, offset, length] + body = section_bytes(section) + return if body.nil? + + if offset + acc["BODY[#{section}]<#{offset}>"] = body.byteslice(offset, length) || +"" + else + acc["BODY[#{section}]"] = body + end + end + + def section_bytes(section) + return "#{@mail.header.raw_source}\r\n" if section == "HEADER" + + locate(section)&.body&.raw_source + end + + # RFC 3501 part addressing: a single-part message is section 1; otherwise walk the + # dotted index path through the multipart tree. + def locate(section) + indices = section.split(".").map { |number| number.to_i - 1 } + return indices == [0] ? @mail : nil unless @mail.multipart? + + indices.reduce(@mail) do |node, index| + child = descend(node, index) + return nil if child.nil? + + child + end + end + + def descend(node, index) + return nil unless node.respond_to?(:multipart?) && node.multipart? + + node.parts[index] + end + + def structure_of(part) + return leaf_of(part) unless part.multipart? + + Net::IMAP::BodyTypeMultipart.new( + "MULTIPART", part.mime_type.to_s.split("/").last.to_s.upcase, + part.parts.map { |child| structure_of(child) }, + params_of(part), disposition_of(part), nil, nil, nil + ) + end + + def leaf_of(part) + media, sub = (part.mime_type || "text/plain").split("/") + size = part.body.raw_source.bytesize + shared = [media.to_s.upcase, sub.to_s.upcase, params_of(part), part.content_id, nil, + (part.content_transfer_encoding || "7bit").to_s, size] + if media.to_s.casecmp?("text") + Net::IMAP::BodyTypeText.new(*shared, part.body.raw_source.count("\n"), nil, + disposition_of(part), nil, nil, nil) + else + Net::IMAP::BodyTypeBasic.new(*shared, nil, disposition_of(part), nil, nil, nil) + end + end + + # Upcased on purpose: real servers pick their own case for parameter names. + def params_of(part) + (part.content_type_parameters || {}).transform_keys { |key| key.to_s.upcase } + rescue StandardError + {} + end + + def disposition_of(part) + dsp = part.header[:content_disposition] + return nil unless dsp + + Net::IMAP::ContentDisposition.new( + dsp.disposition_type.to_s.upcase, + (dsp.parameters || {}).transform_keys { |key| key.to_s.upcase } + ) + rescue StandardError + nil + end +end From 26bd1db6dd46d536032753d4452fa1dd257d8224 Mon Sep 17 00:00:00 2001 From: Ivars Belovs Date: Mon, 24 Aug 2026 14:54:09 +0300 Subject: [PATCH 2/3] Clean up the per-part fetch path Applied findings from a four-angle review of the previous commit. Reuse: net-imap already reconciles the request spec against the response key, which differ (servers answer BODY.PEEK[...] with BODY[...] and append the origin octet on a partial fetch). FetchStruct#part/#header do this, so the hand-rolled body_of prefix scan is gone. Its scan matched any origin, so a server echoing a different one would have been accepted and the loop would have advanced over the wrong bytes; #part matches exactly. A missing section is now distinguishable from end-of-part -- empty means EOF, nil means the server answered something we did not ask for -- so a mid-stream gap raises IncompletePart instead of uploading a zero-byte attachment under a valid-looking presigned URL. The spec fake returns a real Net::IMAP::FetchData so specs exercise that reconciliation rather than a stand-in for it. Correctness of filenames: BODYSTRUCTURE parameters arrive RFC 2231-split across numbered keys (FILENAME*0*, FILENAME*1*) or RFC 2047 encoded. The old casecmp?-based lookup returned nil for both, and since a filename is what marks a part as an attachment, such attachments vanished from the response entirely. Mail::ParameterHash plus Encodings.value_decode handles all three forms. The comment claiming servers vary parameter-name case was also wrong: net-imap's parser upcases them. Efficiency: Base64Stream#push made three full-size copies of every chunk -- measured 15.00 MiB of garbage per 4 MiB chunk. It now strips and slices in place, measured at 3.99 MiB. This required pairing with the caller, which must read chunk.bytesize before handing the buffer over; the chunk fetch and write also moved into their own method so the buffer is unreachable rather than held in a loop local across the next round trip. Rebuilding an inline part now uses setters instead of interpolating a MIME header and re-parsing it, saving about one copy of each inline body. Five concurrent 25MB-attachment fetches now peak at 221MB, against 276MB before this cleanup and 614MB before the streaming change. Simplification: AttachmentStore.upload had no caller left and preserved the fully-resident String upload this work exists to remove -- deleted, with store inlined into upload_io. Part collapsed media_type/subtype into one mime_type member and dropped the never-read content_id. streamable? no longer conflates "we can decode this incrementally" with "the server reported a size"; the chunk loop terminates on the first empty range, so a part of unreported size streams too. Tempfile.create's block form replaces the manual ensure/close!, which also closes a leak: the tempfile used to be created outside the begin, so a fetch raising mid-stream held an fd and its disk until finalization. Also trimmed dead guards in Base64Stream, a guard clause in flatten, and dead code in the fake. The comment on the memoized S3 client claimed ~11MB per client; measured it is ~19MB and ~110ms once, then ~0.8ms -- corrected. Deliberately not changed: - Streaming quoted-printable. Mail's QP decoder normalizes line endings and repairs hard breaks mis-encoded as hex, over the whole string, and is empirically not decomposable across chunk boundaries -- it mismatches at every slice size tested, including when only the trailing 8 bytes are split off. A hand-rolled substitute risks silently corrupting attachment content, which is worse than the bounded-but-larger memory use of an encoding only chosen for text. The one-pass path now logs when taken, and the comment no longer implies large QP parts cannot occur. - Capping inline body size. Inline parts are still fetched whole, so a very large inline HTML body is unbounded. Pre-existing, and truncating a body changes what the tool returns, so it belongs in its own change. - Merging the two inline part fetches into one round trip. Saves one of four fetches but holds both raw bodies at once, on the axis this work exists to bound. - Moving upload and hash assembly out of MessageReader. Offloading attachments is inherent to reading a message without buffering it, and splitting it back out pushes ImapClient over its length limit again or needs a third class. Co-Authored-By: Claude Opus 5 (1M context) --- lib/mail_mcp/attachment_store.rb | 35 +++---- lib/mail_mcp/base64_stream.rb | 17 ++-- lib/mail_mcp/message_reader.rb | 116 +++++++++++++----------- lib/mail_mcp/message_structure.rb | 58 ++++++------ spec/mail_mcp/attachment_store_spec.rb | 42 ++++----- spec/mail_mcp/imap_client_spec.rb | 12 +++ spec/mail_mcp/message_structure_spec.rb | 40 +++++++- spec/support/fake_imap_server.rb | 23 +---- 8 files changed, 187 insertions(+), 156 deletions(-) diff --git a/lib/mail_mcp/attachment_store.rb b/lib/mail_mcp/attachment_store.rb index 24c96d2..d69ccb0 100644 --- a/lib/mail_mcp/attachment_store.rb +++ b/lib/mail_mcp/attachment_store.rb @@ -6,44 +6,35 @@ module AttachmentStore EXPIRY = 7 * 24 * 3600 class << self - # Uploads an in-memory body. Prefer #upload_io for anything attachment-sized: - # a String body has to be fully resident, which is what we are trying to avoid. - def upload(content:, filename:, content_type:) - store(body: content, filename: filename, content_type: content_type) - end - # Uploads from an open IO (typically a Tempfile), letting the SDK read it in # chunks instead of holding the whole attachment as a String. def upload_io(io:, filename:, content_type:) io.rewind - store(body: io, filename: filename, content_type: content_type) - end - - # Drops the memoized clients. Only needed by tests and after credential changes. - def reset! - @s3 = nil - @presigner = nil - end - - private - - def store(body:, filename:, content_type:) key = "attachments/#{SecureRandom.uuid}/#{filename}" bucket = ENV.fetch("AWS_S3_BUCKET") s3.put_object( bucket: bucket, key: key, - body: body, + body: io, content_type: content_type ) presigner.presigned_url(:get_object, bucket: bucket, key: key, expires_in: EXPIRY) end - # Memoized: building a client cost ~11MB of RSS, and the un-memoized version - # built two per attachment — once here and once more inside #presigner. - # Aws::S3::Client is thread-safe, so one instance serves every Puma thread. + # Drops the memoized clients. Only needed by tests and after credential changes. + def reset! + @s3 = nil + @presigner = nil + end + + private + + # Memoized because the first Aws::S3::Client.new costs ~110ms and ~19MB of RSS + # loading plugins, and the un-memoized version built two per attachment — once + # here and once more inside #presigner. Aws::S3::Client is thread-safe, so one + # instance serves every Puma thread. def s3 @s3 ||= Aws::S3::Client.new end diff --git a/lib/mail_mcp/base64_stream.rb b/lib/mail_mcp/base64_stream.rb index b565d80..3fec315 100644 --- a/lib/mail_mcp/base64_stream.rb +++ b/lib/mail_mcp/base64_stream.rb @@ -10,19 +10,20 @@ def initialize @carry = +"" end + # Consumes +chunk+: it is stripped and sliced in place so that decoding a 4MiB + # chunk does not allocate three more copies of it. Callers must read anything they + # need from the chunk (its size, in particular) before calling this. def push(chunk) - buffer = @carry + chunk.delete("\r\n\t ") - complete = buffer.bytesize - (buffer.bytesize % GROUP) - @carry = buffer.byteslice(complete, buffer.bytesize - complete) || +"" - return "".b if complete.zero? - - buffer.byteslice(0, complete).unpack1("m") + chunk = chunk.dup if chunk.frozen? + chunk.delete!("\r\n\t ") + chunk.prepend(@carry) unless @carry.empty? + leftover = chunk.bytesize % GROUP + @carry = leftover.zero? ? +"" : chunk.slice!(chunk.bytesize - leftover, leftover) + chunk.empty? ? "".b : chunk.unpack1("m") end # Decodes whatever is left over, including any padding. def finish - return "".b if @carry.empty? - remainder = @carry @carry = +"" remainder.unpack1("m") diff --git a/lib/mail_mcp/message_reader.rb b/lib/mail_mcp/message_reader.rb index 8e86841..7222e65 100644 --- a/lib/mail_mcp/message_reader.rb +++ b/lib/mail_mcp/message_reader.rb @@ -5,6 +5,10 @@ module MailMCP # tracks the part being handled rather than the size of the whole message. Fetching # RFC822 instead would materialize the entire message as a single String. class MessageReader + # Raised when the server answers a section fetch with a section we did not ask for. + # Treated as an error because the alternative is a silently truncated attachment. + class IncompletePart < StandardError; end + # Encoded bytes pulled per FETCH when streaming an attachment. Bounds peak memory # regardless of attachment size; larger chunks mean fewer round trips but hold a # Puma thread for longer. @@ -23,26 +27,12 @@ def read parts = MessageStructure.flatten(data.attr["BODYSTRUCTURE"]) MailMCP.logger.debug { "IMAP message uid=#{@uid} parts=#{parts.map(&:section).inspect}" } - format( - headers: Mail.new(self.class.body_of(data.attr, "HEADER").to_s), - flags: data.attr["FLAGS"], - parts: parts - ) - end - - # Servers answer BODY.PEEK[...] with BODY[...], and append the origin octet on a - # partial fetch (BODY[1]<0>), so the response key never matches the request. - def self.body_of(attrs, section) - exact = attrs["BODY[#{section}]"] - return exact if exact - - prefix = "BODY[#{section}]<" - attrs.find { |key, _value| key.to_s.start_with?(prefix) }&.last + to_h(headers: Mail.new(data.header.to_s), flags: data.attr["FLAGS"], parts: parts) end private - def format(headers:, flags:, parts:) + def to_h(headers:, flags:, parts:) inline = parts.reject(&:attachment?) { uid: @uid, @@ -74,14 +64,15 @@ def inline_body(inline, mime_type) body && decode_inline(part, body) end - # Rebuilding the part from a synthesized MIME header lets Mail apply both the - # transfer-encoding and the charset conversion, which is what #decoded did on the - # old whole-message parse. Hand-rolling that would mislabel any non-UTF-8 body. + # Letting Mail decode a rebuilt part applies both the transfer-encoding and the + # charset conversion, which is what #decoded did on the old whole-message parse. + # Hand-rolling that would mislabel any non-UTF-8 body. def decode_inline(part, body) - header = "Content-Type: #{part.mime_type}" - header += "; charset=#{part.charset}" if part.charset - header += "\r\nContent-Transfer-Encoding: #{part.encoding}\r\n" - Mail::Part.new("#{header}\r\n#{body}").decoded + rebuilt = Mail::Part.new + rebuilt.content_type = part.charset ? "#{part.mime_type}; charset=#{part.charset}" : part.mime_type + rebuilt.content_transfer_encoding = part.encoding + rebuilt.body = body + rebuilt.decoded rescue StandardError => e MailMCP.logger.warn { "IMAP inline decode failed section=#{part.section}: #{e.class}: #{e.message}" } nil @@ -91,47 +82,58 @@ def decode_inline(part, body) # fetched, so attachments never accumulate in memory. def upload_attachments(parts) parts.map do |part| - tempfile = fetch_attachment(part) - begin + Tempfile.create("mail_mcp_attachment", binmode: true) do |io| + size = fetch_attachment(io, part) url = AttachmentStore.upload_io( - io: tempfile, + io: io, filename: part.filename || "attachment", content_type: part.mime_type ) - { filename: part.filename, content_type: part.mime_type, size: tempfile.size, url: url } - ensure - tempfile.close! + { filename: part.filename, content_type: part.mime_type, size: size, url: url } end end end - def fetch_attachment(part) - tempfile = Tempfile.new("mail_mcp_attachment") - tempfile.binmode - if part.streamable? - stream_section(tempfile, part) - else - # Quoted-printable, or a part whose size the server did not report. Neither is - # used for large payloads in practice, so one pass is acceptable here. - tempfile.write(decode_whole(fetch_section(part.section).to_s, part.encoding)) + # Returns the decoded byte count written to +io+. + def fetch_attachment(io, part) + return stream_section(io, part) if part.streamable? + + # Quoted-printable has to be decoded in one pass, so this path is bounded by the + # part size rather than by CHUNK_SIZE. Logged because it is the one place where + # the memory guarantee above does not hold. + MailMCP.logger.warn do + "IMAP attachment decoded in one pass encoding=#{part.encoding.inspect} " \ + "encoded_size=#{part.encoded_size} section=#{part.section}" end - tempfile.flush - tempfile + io.write(decode_whole(require_section(part.section), part.encoding)) end - def stream_section(tempfile, part) + def stream_section(io, part) decoder = part.base64? ? Base64Stream.new : nil offset = 0 - # Bounded by the size BODYSTRUCTURE reported; the empty check is a second guard so - # a server returning short reads cannot spin here. - while offset < part.encoded_size - chunk = fetch_section(part.section, offset: offset, length: CHUNK_SIZE) - break if chunk.nil? || chunk.empty? - - tempfile.write(decoder ? decoder.push(chunk) : chunk) - offset += chunk.bytesize + written = 0 + loop do + fetched, wrote = consume_chunk(io, decoder, part, offset) + break if fetched.zero? + + offset += fetched + written += wrote end - tempfile.write(decoder.finish) if decoder + written += io.write(decoder.finish) if decoder + written + end + + # Kept in its own method so the chunk becomes unreachable as soon as it returns. + # Held in a local in the loop above, a CHUNK_SIZE buffer would stay alive across + # the next fetch, doubling the resident cost of streaming. + # Returns [encoded bytes fetched, decoded bytes written]. + def consume_chunk(io, decoder, part, offset) + chunk = require_section(part.section, offset: offset) + return [0, 0] if chunk.empty? + + # #push consumes the chunk, so read its size before handing it over. + fetched = chunk.bytesize + [fetched, io.write(decoder ? decoder.push(chunk) : chunk)] end def decode_whole(raw, encoding) @@ -142,10 +144,18 @@ def decode_whole(raw, encoding) raw end - def fetch_section(section, offset: nil, length: nil) - spec = offset ? "BODY.PEEK[#{section}]<#{offset}.#{length}>" : "BODY.PEEK[#{section}]" + def require_section(section, offset: nil) + fetch_section(section, offset: offset) || + raise(IncompletePart, "IMAP returned no BODY[#{section}] at offset #{offset.inspect} for uid #{@uid}") + end + + # net-imap reconciles the request spec against the response key, which differ: + # servers answer BODY.PEEK[...] with BODY[...] and append the origin octet on a + # partial fetch. A nil result means a section we did not ask for. + def fetch_section(section, offset: nil) + spec = offset ? "BODY.PEEK[#{section}]<#{offset}.#{CHUNK_SIZE}>" : "BODY.PEEK[#{section}]" data = @imap.uid_fetch([@uid], [spec])&.first - data && self.class.body_of(data.attr, section) + data&.part(*section.split("."), offset: offset) end end end diff --git a/lib/mail_mcp/message_structure.rb b/lib/mail_mcp/message_structure.rb index aaedac7..d26b25a 100644 --- a/lib/mail_mcp/message_structure.rb +++ b/lib/mail_mcp/message_structure.rb @@ -1,4 +1,4 @@ -require "net/imap" +require "mail" module MailMCP # Flattens an IMAP BODYSTRUCTURE into its leaf parts, each tagged with the @@ -6,14 +6,13 @@ module MailMCP # individually is what lets a large message be handled without ever holding the # whole thing in memory. module MessageStructure - IDENTITY_ENCODINGS = ["7bit", "8bit", "binary", ""].freeze - - Part = Struct.new(:section, :media_type, :subtype, :encoding, :encoded_size, :filename, - :charset, :content_id, keyword_init: true) do - def mime_type - "#{media_type}/#{subtype}" - end + # Encodings we can fetch and write through in chunks. Quoted-printable cannot be: + # Mail's decoder normalizes line endings and repairs mis-encoded hard breaks across + # the whole string, and does not decompose across chunk boundaries. + STREAMABLE_ENCODINGS = ["base64", "7bit", "8bit", "binary", ""].freeze + Part = Struct.new(:section, :mime_type, :encoding, :encoded_size, :filename, :charset, + keyword_init: true) do # Mirrors Mail::Message#attachment?, which keys off a filename being present # rather than off the Content-Disposition type. def attachment? @@ -24,23 +23,18 @@ def base64? encoding == "base64" end - # Encodings we can fetch and write through in chunks. Anything else (in - # practice quoted-printable) has to be decoded in one pass. def streamable? - encoded_size.positive? && (base64? || IDENTITY_ENCODINGS.include?(encoding)) + STREAMABLE_ENCODINGS.include?(encoding) end end class << self def flatten(body, prefix = nil) return [] if body.nil? + return [leaf(body, prefix || "1")] unless body.multipart? - if body.multipart? - body.parts.each_with_index.flat_map do |child, index| - flatten(child, [prefix, index + 1].compact.join(".")) - end - else - [leaf(body, prefix || "1")] + body.parts.flat_map.with_index(1) do |child, number| + flatten(child, [prefix, number].compact.join(".")) end end @@ -49,26 +43,30 @@ def flatten(body, prefix = nil) def leaf(body, section) Part.new( section: section, - media_type: body.media_type.to_s.downcase, - subtype: body.subtype.to_s.downcase, + mime_type: "#{body.media_type}/#{body.subtype}".downcase, encoding: body.encoding.to_s.downcase, encoded_size: body.size.to_i, - filename: filename_for(body), - charset: param(body.param, "charset"), - content_id: body.content_id + filename: param(body.disposition&.param, "filename") || param(body.param, "name"), + charset: param(body.param, "charset") ) end - def filename_for(body) - param(body.disposition&.param, "filename") || param(body.param, "name") - end - - # Servers pick their own case for parameter names, so never index directly. + # A parameter can arrive RFC 2231-split across numbered keys (FILENAME*0*, + # FILENAME*1*) or RFC 2047 encoded. Mail decodes both; a plain key lookup returns + # nil for the split form, which silently drops the attachment from the response. def param(params, name) - return nil unless params + return nil if params.nil? || params.empty? + + value = Mail::ParameterHash[params][name] + return nil if value.nil? + + # An RFC 2231 extended value is prefixed with charset'language'. + value = value.sub(/\A[\w-]*'[\w-]*'/, "") if extended?(params, name) + Mail::Encodings.value_decode(value) + end - pair = params.find { |key, _value| key.to_s.casecmp?(name) } - pair&.last + def extended?(params, name) + params.any? { |key, _value| key.to_s.downcase.start_with?("#{name}*") } end end end diff --git a/spec/mail_mcp/attachment_store_spec.rb b/spec/mail_mcp/attachment_store_spec.rb index 880f928..986ad03 100644 --- a/spec/mail_mcp/attachment_store_spec.rb +++ b/spec/mail_mcp/attachment_store_spec.rb @@ -15,10 +15,10 @@ described_class.reset! end - describe ".upload" do + describe ".upload_io" do it "uploads to S3 and returns a presigned URL" do - url = described_class.upload( - content: "PDF content", + url = described_class.upload_io( + io: StringIO.new("PDF content"), filename: "report.pdf", content_type: "application/pdf" ) @@ -28,8 +28,23 @@ expect(url).to eq(presigned_url) end + # Passing the IO through rather than its contents is what keeps a large attachment + # off the heap, so assert the body is the handle itself. + it "streams the IO to S3 as the request body" do + io = StringIO.new("streamed bytes") + described_class.upload_io(io: io, filename: "big.bin", content_type: "application/octet-stream") + expect(s3_client).to have_received(:put_object).with(hash_including(body: io)) + end + + it "rewinds the IO so a partially read tempfile still uploads in full" do + io = StringIO.new("streamed bytes") + io.read(4) + described_class.upload_io(io: io, filename: "big.bin", content_type: "application/octet-stream") + expect(io.pos).to eq(0) + end + it "reuses one memoized client and presigner across uploads" do - 2.times { described_class.upload(content: "d", filename: "f.txt", content_type: "text/plain") } + 2.times { described_class.upload_io(io: StringIO.new("d"), filename: "f.txt", content_type: "text/plain") } expect(Aws::S3::Client).to have_received(:new).once expect(Aws::S3::Presigner).to have_received(:new).once end @@ -38,27 +53,10 @@ keys = 2.times.map do key = nil allow(s3_client).to receive(:put_object) { |args| key = args[:key] } - described_class.upload(content: "data", filename: "file.txt", content_type: "text/plain") + described_class.upload_io(io: StringIO.new("data"), filename: "file.txt", content_type: "text/plain") key end expect(keys.first).not_to eq(keys.last) end end - - describe ".upload_io" do - it "streams the IO to S3 as the request body" do - io = StringIO.new("streamed bytes") - url = described_class.upload_io(io: io, filename: "big.bin", content_type: "application/octet-stream") - - expect(s3_client).to have_received(:put_object).with(hash_including(body: io)) - expect(url).to eq(presigned_url) - end - - it "rewinds the IO so a partially read tempfile still uploads in full" do - io = StringIO.new("streamed bytes") - io.read(4) - described_class.upload_io(io: io, filename: "big.bin", content_type: "application/octet-stream") - expect(io.pos).to eq(0) - end - end end diff --git a/spec/mail_mcp/imap_client_spec.rb b/spec/mail_mcp/imap_client_spec.rb index 7e737b5..8caf59b 100644 --- a/spec/mail_mcp/imap_client_spec.rb +++ b/spec/mail_mcp/imap_client_spec.rb @@ -110,6 +110,18 @@ expect(body_specs).to all(start_with("BODY.PEEK[")) end + # A missing section mid-stream used to produce a zero-byte attachment with a + # valid-looking presigned URL. Failing loudly beats uploading a truncated file. + it "raises rather than uploading a truncated attachment when a section is missing" do + allow(server).to receive(:uid_fetch).and_wrap_original do |original, uids, specs| + specs.any? { |spec| spec.start_with?("BODY.PEEK[2]") } ? [] : original.call(uids, specs) + end + + expect { described_class.new(server).get_message(folder: "INBOX", uid: 42) } + .to raise_error(MailMCP::MessageReader::IncompletePart, /BODY\[2\]/) + expect(uploaded).to be_empty + end + it "returns nil when the uid is not found" do allow(server).to receive(:uid_fetch).and_return([]) expect(described_class.new(server).get_message(folder: "INBOX", uid: 99)).to be_nil diff --git a/spec/mail_mcp/message_structure_spec.rb b/spec/mail_mcp/message_structure_spec.rb index a0d7ecc..1fcab18 100644 --- a/spec/mail_mcp/message_structure_spec.rb +++ b/spec/mail_mcp/message_structure_spec.rb @@ -68,6 +68,36 @@ def parts_for(raw) expect(parts.select(&:attachment?).map(&:filename)).to eq(["invoice.pdf"]) end + # net-imap hands back RFC 2231 continuations as separate numbered keys. A plain + # key lookup misses them, and the attachment then vanishes from the response + # entirely, since a filename is what marks a part as an attachment. + it "reassembles an RFC 2231 split filename" do + params = { "FILENAME*0*" => "utf-8%27%27r%C3%A4ch", "FILENAME*1*" => "nung%2Epdf" } + disposition = Net::IMAP::ContentDisposition.new("ATTACHMENT", params) + body = Net::IMAP::BodyTypeBasic.new("APPLICATION", "PDF", nil, nil, nil, "base64", 10, + nil, disposition, nil, nil, nil) + + expect(described_class.flatten(body).first.filename).to eq("rächnung.pdf") + end + + it "decodes an RFC 2047 encoded filename" do + disposition = Net::IMAP::ContentDisposition.new( + "ATTACHMENT", { "FILENAME" => "=?UTF-8?B?csOkY2hudW5nLnBkZg==?=" } + ) + body = Net::IMAP::BodyTypeBasic.new("APPLICATION", "PDF", nil, nil, nil, "base64", 10, + nil, disposition, nil, nil, nil) + + expect(described_class.flatten(body).first.filename).to eq("rächnung.pdf") + end + + it "leaves an apostrophe in an ordinary filename alone" do + disposition = Net::IMAP::ContentDisposition.new("ATTACHMENT", { "FILENAME" => "'quoted'.pdf" }) + body = Net::IMAP::BodyTypeBasic.new("APPLICATION", "PDF", nil, nil, nil, "base64", 10, + nil, disposition, nil, nil, nil) + + expect(described_class.flatten(body).first.filename).to eq("'quoted'.pdf") + end + it "reads a filename that the server reported in upper case" do mail = Mail.new do from "a@x.com" @@ -94,9 +124,15 @@ def parts_for(raw) expect(part).not_to be_streamable end - # Without a size from the server the chunk loop has no bound to walk. - it "does not stream a part of unreported size" do + # A part whose size the server did not report still streams: the chunk loop + # terminates on the first empty range rather than on a byte count. + it "streams a part of unreported size" do part = described_class::Part.new(section: "1", encoding: "base64", encoded_size: 0) + expect(part).to be_streamable + end + + it "does not stream an encoding it has no incremental decoder for" do + part = described_class::Part.new(section: "1", encoding: "x-uuencode", encoded_size: 10) expect(part).not_to be_streamable end end diff --git a/spec/support/fake_imap_server.rb b/spec/support/fake_imap_server.rb index ce59c87..70061aa 100644 --- a/spec/support/fake_imap_server.rb +++ b/spec/support/fake_imap_server.rb @@ -14,14 +14,14 @@ def initialize(raw) def examine(_folder) = nil - def select(_folder) = nil - def uid_fetch(_uids, specs) @requested_specs.concat(specs) attrs = specs.each_with_object({}) { |spec, acc| serve(spec, acc) } return nil if attrs.empty? - [Struct.new(:attr).new(attrs)] + # A real FetchData, so specs exercise net-imap's own request-spec/response-key + # reconciliation rather than a stand-in for it. + [Net::IMAP::FetchData.new(1, attrs)] end # Sections fetched with a byte range, e.g. [["2", 0, 4194304]]. @@ -64,18 +64,7 @@ def locate(section) indices = section.split(".").map { |number| number.to_i - 1 } return indices == [0] ? @mail : nil unless @mail.multipart? - indices.reduce(@mail) do |node, index| - child = descend(node, index) - return nil if child.nil? - - child - end - end - - def descend(node, index) - return nil unless node.respond_to?(:multipart?) && node.multipart? - - node.parts[index] + indices.reduce(@mail) { |node, index| node.parts[index] if node&.multipart? } end def structure_of(part) @@ -104,8 +93,6 @@ def leaf_of(part) # Upcased on purpose: real servers pick their own case for parameter names. def params_of(part) (part.content_type_parameters || {}).transform_keys { |key| key.to_s.upcase } - rescue StandardError - {} end def disposition_of(part) @@ -116,7 +103,5 @@ def disposition_of(part) dsp.disposition_type.to_s.upcase, (dsp.parameters || {}).transform_keys { |key| key.to_s.upcase } ) - rescue StandardError - nil end end From 7d3d12ebe305cd56d97435807c0024dcfd062bd2 Mon Sep 17 00:00:00 2001 From: Ivars Belovs Date: Mon, 24 Aug 2026 15:20:26 +0300 Subject: [PATCH 3/3] Pin the attached-multipart flattening semantics Copilot flagged flatten's unconditional descent into multipart nodes as dropping an attached multipart and letting its children supply text_body/html_body. Checked against the old whole-message path with a multipart/related carrying Content-Disposition: attachment; filename="page.mht". The new output matches the old one exactly on all three fields: no attachment, text_body from the real body part, html_body from inside the archive. That is Mail's own behaviour, not an accident. Mail::AttachmentsList (attachments_list.rb:8-19) only ever treats leaf parts as attachments -- a part with children recurses and is never included itself, whatever its disposition says. message/rfc822 is its one special case, and that already arrives here as a leaf because BodyTypeMessage#multipart? is false, so it is picked up as an attachment as before. So there is no regression to fix, and emitting the multipart node as an attachment would both diverge from Mail and add an attachment to responses the old path never returned. Added a spec so the semantics are deliberate rather than incidental. The second half of the observation is fair on its own terms: html_body can come from inside an attached archive. That predates this branch, matches Mail#html_part, and changing it changes what the tool returns, so it belongs in its own change. Co-Authored-By: Claude Opus 5 (1M context) --- spec/mail_mcp/message_structure_spec.rb | 37 +++++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/spec/mail_mcp/message_structure_spec.rb b/spec/mail_mcp/message_structure_spec.rb index 1fcab18..9d38c39 100644 --- a/spec/mail_mcp/message_structure_spec.rb +++ b/spec/mail_mcp/message_structure_spec.rb @@ -10,6 +10,20 @@ def parts_for(raw) described_class.flatten(structure_for(raw)) end + def html_part(body) + part = Mail::Part.new + part.content_type = "text/html" + part.body = body + part + end + + def png_part(body) + part = Mail::Part.new + part.content_type = "image/png" + part.body = body + part + end + it "returns nothing for a nil structure" do expect(described_class.flatten(nil)).to eq([]) end @@ -54,6 +68,29 @@ def parts_for(raw) expect(parts_for(outer.to_s).map(&:section)).to eq(%w[1.1 1.2 2]) end + # Mail::AttachmentsList only ever treats leaf parts as attachments: a node with + # children recurses and is never included itself, whatever its disposition says + # (message/rfc822 is its one special case, and that arrives here as a leaf because + # BodyTypeMessage#multipart? is false). Descending unconditionally reproduces that. + it "descends into a multipart that carries its own attachment disposition" do + inner = Mail::Part.new + inner.content_type = "multipart/related; boundary=INNER" + inner.content_disposition = 'attachment; filename="page.mht"' + inner.add_part(html_part("

archived

")) + inner.add_part(png_part("PNGDATA")) + + outer = Mail.new + outer.from = "a@x.com" + outer.to = "b@x.com" + outer.text_part = Mail::Part.new { body "real body" } + outer.add_part(inner) + + parts = parts_for(outer.to_s) + expect(parts.map(&:mime_type)).to eq(["text/plain", "text/html", "image/png"]) + expect(parts.map(&:section)).to eq(%w[1 2.1 2.2]) + expect(parts).to all(satisfy { |part| part.filename.nil? }) + end + describe "attachment detection" do it "treats a part with a filename as an attachment" do mail = Mail.new do