EMAIL-204: Disable eager attachment loading on MimeMessageParser

src/main/java/org/apache/commons/mail/LazyByteArrayDataSource.java

@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.commons.mail;
+
+import javax.activation.DataSource;
+import javax.mail.util.ByteArrayDataSource;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+
+/**
+ * <p>Wrapper class for ByteArrayDataSource, which contain reference of MimePartDataSource for given attachment.
+ * Both type and name are duplicated stored in this class, in order to delay the load of attachment binary till getInputStream() is called.
+ * </p>
+ *
+ * @since 1.5
+ */
+public class LazyByteArrayDataSource implements DataSource {
+
+    /** InputStream reference for the email attachment binary. */
+    private final InputStream referenceInputStream;
+
+    /** ByteArrayDateSource instance which contain email attachment binary in the form of byte array. */
+    private ByteArrayDataSource ds;
+
+    /** Name of the attachment. */
+    private final String name;
+
+    /** Type of the attachment. */
+    private final String type;
+
+
+    /**
+     * Constructor for this class to read all necessary information for an email attachment.
+     *
+     * @param is the InputStream which represent the attachment binary.
+     * @param type the type of the attachment.
+     * @param name the name of the attachment.
+     */
+    public LazyByteArrayDataSource(InputStream is, String type, String name) {
+        this.referenceInputStream = is;
+        this.type = type;
+        this.name = name;
+    }
+
+    /**
+     * To return an {@code ByteArrayDataSource} instance which represent the email attachment.
+     *
+     * @return An ByteArrayDataSource instance which contain the email attachment.
+     * @throws IOException resolving the email attachment failed
+     */
+    @Override
+    public InputStream getInputStream() throws IOException {
+        if (ds == null) {
+            //Only read attachment data to memory when getInputStream() is called.
+            ds = new ByteArrayDataSource(referenceInputStream, type);
+            ds.setName(name);
+        }
+        return ds.getInputStream();
+    }
+
+    /**
+     * Not supported.
+     *
+     * @return  N/A
+     * @since 1.5
+     */
+    @Override
+    public OutputStream getOutputStream() throws IOException {
+        throw new IOException("cannot do this");
+    }
+
+    /**
+     * Get the content type.
+     *
+     * @return A String.
+     * @since 1.5
+     */
+    @Override
+    public String getContentType() {
+        return type;
+    }
+
+    /**
+     * Get the name.
+     *
+     * @return A String.
+     * @since 1.5
+     */
+    @Override
+    public String getName() {
+        return name;
+    }
+}

src/main/java/org/apache/commons/mail/util/MimeMessageParser.java

@@ -16,11 +16,9 @@
  */
 package org.apache.commons.mail.util;
 
-import java.io.BufferedInputStream;
-import java.io.BufferedOutputStream;
-import java.io.ByteArrayOutputStream;
+import org.apache.commons.mail.LazyByteArrayDataSource;
+
 import java.io.IOException;
-import java.io.InputStream;
 import java.io.UnsupportedEncodingException;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -43,7 +41,6 @@
 import javax.mail.internet.MimePart;
 import javax.mail.internet.MimeUtility;
 import javax.mail.internet.ParseException;
-import javax.mail.util.ByteArrayDataSource;
 
 /**
  * Parses a MimeMessage and stores the individual parts such a plain text,
@@ -270,14 +267,8 @@ protected DataSource createDataSource(final Multipart parent, final MimePart par
         final DataHandler dataHandler = part.getDataHandler();
         final DataSource dataSource = dataHandler.getDataSource();
         final String contentType = getBaseMimeType(dataSource.getContentType());
-        byte[] content;
-        try (InputStream inputStream = dataSource.getInputStream()) 
-        {
-            content = this.getContent(inputStream);
-        }
-        final ByteArrayDataSource result = new ByteArrayDataSource(content, contentType);
         final String dataSourceName = getDataSourceName(part, dataSource);
-        result.setName(dataSourceName);
+        final LazyByteArrayDataSource result = new LazyByteArrayDataSource(dataSource.getInputStream(), contentType, dataSourceName);
         return result;
     }
 
@@ -411,28 +402,6 @@ protected String getDataSourceName(final Part part, final DataSource dataSource)
         return result;
     }
 
-    /**
-     * Read the content of the input stream.
-     *
-     * @param is the input stream to process
-     * @return the content of the input stream
-     * @throws IOException reading the input stream failed
-     */
-    private byte[] getContent(final InputStream is)
-        throws IOException
-    {
-        final ByteArrayOutputStream os = new ByteArrayOutputStream();
-        final BufferedInputStream isReader = new BufferedInputStream(is);
-        try (BufferedOutputStream osWriter = new BufferedOutputStream(os)) {
-            int ch;
-            while ((ch = isReader.read()) != -1)
-            {
-                osWriter.write(ch);
-            }
-            osWriter.flush();
-            return os.toByteArray();
-        }
-    }
 
     /**
      * Parses the mimeType.

src/test/java/org/apache/commons/mail/util/MimeMessageParserTest.java

@@ -16,19 +16,24 @@
  */
 package org.apache.commons.mail.util;
 
+import static org.easymock.EasyMock.expect;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
+import static org.powermock.api.easymock.PowerMock.*;
 
 import java.io.File;
+import java.io.InputStream;
 import java.util.List;
 import java.util.Properties;
 
+import javax.activation.DataHandler;
 import javax.activation.DataSource;
 import javax.mail.Session;
 import javax.mail.internet.MimeMessage;
+import javax.mail.internet.MimePart;
 
 import org.apache.commons.mail.HtmlEmail;
 import org.junit.Test;
@@ -498,4 +503,61 @@ public void testParseInlineCID() throws Exception
         assertEquals(ds, mimeMessageParser.getAttachmentList().get(0));
     }
 
+    @Test
+    public void testAttachmentNotLoaded() throws Exception
+    {
+        final MimeMessageParser mimeMessageParser = new MimeMessageParser(null);
+        final InputStream inputStream = createMock(InputStream.class);
+        final MimePart mimePart = getMockedMimePart(inputStream);
+
+        // Create data source with mocked data.
+        final DataSource dataSource_new = mimeMessageParser.createDataSource(null,mimePart);
+        // Verify no inputStream.read() is made at this point (Lazy initialization).
+        verify(inputStream);
+    }
+
+
+    @Test
+    public void testAttachmentLoaded() throws Exception
+    {
+        final MimeMessageParser mimeMessageParser = new MimeMessageParser(null);
+        final InputStream inputStream = createMock(InputStream.class);
+        // Despite .getInputStream() called for 3 times, but the desk IO for attachment read should only happen once.
+        expect(inputStream.read(new byte[8192])).andReturn(0).once();
+        final MimePart mimePart = getMockedMimePart(inputStream);
+
+        // Create data source with mocked data.
+        final DataSource dataSource_new = mimeMessageParser.createDataSource(null,mimePart);
+
+        dataSource_new.getInputStream();
+        dataSource_new.getInputStream();
+        dataSource_new.getInputStream();
+        // To make sure disk IO only happen when .getInputStream() invoked for first time but during the object construction.
+        verify(inputStream);
+
+    }
+
+    /**
+     * Helper method to return a mocked MimePart class.
+     * @param inputStream Mocked input stream
+     * @return Mocked MimePart instance.
+     * @throws Exception When attachment read failed.
+     */
+    private MimePart getMockedMimePart(InputStream inputStream) throws Exception
+    {
+
+        final MimePart mimePart = createMock(MimePart.class);
+        final DataHandler dataHandler = createMock(DataHandler.class);
+        final DataSource dataSource = createMock(DataSource.class);
+
+        expect(dataSource.getContentType()).andReturn("test_type");
+        expect(dataSource.getName()).andReturn("test_name");
+        expect(dataSource.getInputStream()).andReturn(inputStream).once();
+        expect(mimePart.getDataHandler()).andReturn(dataHandler);
+        expect(dataHandler.getDataSource()).andReturn(dataSource);
+        replay(mimePart,dataHandler,dataSource,inputStream);
+
+        return mimePart;
+    }
+
 }