replace inline attachments identified with a Content-ID

pom.xml

@@ -6,7 +6,7 @@
 
     <groupId>com.massaro</groupId>
     <artifactId>Email</artifactId>
-    <version>2.3.1</version>
+    <version>2.4</version>
 
     <parent>
         <groupId>org.springframework.boot</groupId>

src/main/java/email/endpoint/AttachmentEndpoint.java

@@ -16,6 +16,7 @@
 import org.springframework.web.bind.annotation.RestController;
 import org.springframework.web.server.ResponseStatusException;
 
+import java.util.Map;
 import java.util.Optional;
 
 @RestController
@@ -30,19 +31,39 @@ public AttachmentEndpoint(MessageService messageService) {
 
     @GetMapping
     public ResponseEntity<Resource> get(@RequestParam("id") long id) {
+        ResponseEntity<Resource> resourceResponseEntity = checkAttachments(id);
+        if (resourceResponseEntity != null) {
+            return resourceResponseEntity;
+        }
+
+        resourceResponseEntity = checkCid(id);
+        if (resourceResponseEntity != null) {
+            return resourceResponseEntity;
+        }
+        throw new ResponseStatusException(HttpStatus.NOT_FOUND, "no message found with the corresponding attachment ID");
+    }
+
+    private ResponseEntity<Resource> checkAttachments(long id) {
         Optional<Message> messageWithAttachment = messageService.list().stream()
                 .filter(m -> m.getAttachments() != null && !m.getAttachments().isEmpty() && m.getAttachments().stream().anyMatch(a -> a.getId() == id)).findFirst();
         if (messageWithAttachment.isPresent()) {
             Attachment attachment = messageWithAttachment.get().getAttachments().stream().filter(a -> a.getId() == id).findFirst().orElse(null);
             if (attachment != null) {
-                Resource file = new ByteArrayResource(attachment.getFile());
-                return ResponseEntity
-                        .ok()
-                        .header(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"" + attachment.getName() + "\"")
-                        .contentType(MediaType.parseMediaType(attachment.getContentType()))
-                        .body(file);
+                return attachment.toResponseEntity();
             }
         }
-        throw new ResponseStatusException(HttpStatus.NOT_FOUND, "no message found with the corresponding attachment ID");
+        return null;
+    }
+
+    private ResponseEntity<Resource> checkCid(long id) {
+        Optional<Message> messageWithAttachment = messageService.list().stream()
+                .filter(m -> m.getCidMap() != null && !m.getCidMap().isEmpty() && m.getCidMap().entrySet().stream().anyMatch(a -> a.getValue().getId() == id)).findFirst();
+        if (messageWithAttachment.isPresent()) {
+            Map.Entry<String, Attachment> attachment = messageWithAttachment.get().getCidMap().entrySet().stream().filter(a -> a.getValue().getId() == id).findFirst().orElse(null);
+            if (attachment != null) {
+                return attachment.getValue().toResponseEntity();
+            }
+        }
+        return null;
     }
 }

src/main/java/email/endpoint/BodyEndpoint.java

@@ -1,5 +1,6 @@
 package email.endpoint;
 
+import email.model.Attachment;
 import email.model.BodyPart;
 import email.model.ContentTypeEnum;
 import email.model.Message;
@@ -52,6 +53,14 @@ public ResponseEntity<String> getBody(@RequestParam("id") long id) throws Messag
             }
         }
 
+        // replace all cid (Content-Id) identifiers with a link to the attachment endpoint
+        for (Map.Entry<String, Attachment> cidMapEntry : message.getCidMap().entrySet()) {
+            Attachment attachment = cidMapEntry.getValue();
+            String contentId = cidMapEntry.getKey();
+
+            body = body.replace("cid:" + contentId, "./attachment?id=" + attachment.getId());
+        }
+
         return new ResponseEntity<>(body, responseHeaders, HttpStatus.OK);
     }
 }

src/main/java/email/model/Attachment.java

@@ -4,6 +4,16 @@
 import email.service.MessageService;
 import lombok.Getter;
 import lombok.Setter;
+import org.apache.commons.io.IOUtils;
+import org.apache.commons.lang3.StringUtils;
+import org.springframework.core.io.ByteArrayResource;
+import org.springframework.core.io.Resource;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.MediaType;
+import org.springframework.http.ResponseEntity;
+
+import javax.activation.DataSource;
+import java.io.IOException;
 
 @Getter
 @Setter
@@ -23,4 +33,27 @@ public Attachment(String name, String contentType, byte[] file) {
         this.contentType = contentType;
         this.file = file;
     }
+
+    public ResponseEntity<Resource> toResponseEntity() {
+        Resource file = new ByteArrayResource(getFile());
+        return ResponseEntity
+                .ok()
+                .header(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"" + getName() + "\"")
+                .contentType(MediaType.parseMediaType(getContentType()))
+                .body(file);
+    }
+
+    /**
+     * Given a {@link DataSource}, create an {@link Attachment} object from it.
+     * @return the attachment, or null if the DataSource is missing required parts (name, contentType, data).
+     */
+    public static Attachment fromDataSource(DataSource dataSource) throws IOException {
+        String name = dataSource.getName();
+        String contentType = dataSource.getContentType();
+        byte[] file = IOUtils.toByteArray(dataSource.getInputStream());
+        if (StringUtils.isNotEmpty(name) && StringUtils.isNotEmpty(contentType) && file != null) {
+            return new Attachment(dataSource.getName(), dataSource.getContentType(), IOUtils.toByteArray(dataSource.getInputStream()));
+        }
+        return null;
+    }
 }

src/main/java/email/model/Message.java

@@ -19,9 +19,7 @@
 import javax.mail.internet.MimeMessage;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
-import java.util.ArrayList;
-import java.util.Date;
-import java.util.List;
+import java.util.*;
 
 import static email.model.ProviderEnum.AOL;
 
@@ -44,6 +42,10 @@ public class Message {
     private List<BodyPart> bodyParts = new ArrayList<>();
     private boolean readInd;
     private List<Attachment> attachments = new ArrayList<>();
+    /**
+     * Map of Content-ID's to their respective attachments.
+     */
+    private Map<String, Attachment> cidMap = new HashMap<>();
 
     public Message() {
 
@@ -58,6 +60,7 @@ public Message(javax.mail.Message message, long uid, boolean alreadyExists) thro
                 mimeMessageParser.parse();
                 setAttachments(mimeMessageParser);
                 setBodyParts(mimeMessageParser);
+                setCidMap(mimeMessageParser);
             } catch (Exception e) {
                 log.warn("Failed to parse message using commons email parser, resorting to fallback method (which is less mature)", e);
                 setBodyParts(message);
@@ -76,6 +79,17 @@ public Message(javax.mail.Message message, long uid, boolean alreadyExists) thro
         this.readInd = determineReadInd(message);
     }
 
+    private void setCidMap(MimeMessageParser mimeMessageParser) throws IOException {
+        Collection<String> contentIds = mimeMessageParser.getContentIds();
+        for (String contentId : contentIds) {
+            DataSource attachmentDs = mimeMessageParser.findAttachmentByCid(contentId);
+            Attachment attachment = Attachment.fromDataSource(attachmentDs);
+            if (attachment != null) {
+                this.cidMap.put(contentId, attachment);
+            }
+        }
+    }
+
     public static Date getReceivedDate(javax.mail.Message message) throws MessagingException {
         try {
             String[] dates = message.getHeader("Date");
@@ -158,12 +172,10 @@ public void setBodyParts(MimeMessageParser mimeMessageParser) {
 
     public void setAttachments(MimeMessageParser mimeMessageParser) throws IOException {
         List<DataSource> attachments = mimeMessageParser.getAttachmentList();
-        for (DataSource attachment : attachments) {
-            String name = attachment.getName();
-            String contentType = attachment.getContentType();
-            byte[] file = IOUtils.toByteArray(attachment.getInputStream());
-            if (StringUtils.isNotEmpty(name) && StringUtils.isNotEmpty(contentType) && file != null) {
-                this.attachments.add(new Attachment(attachment.getName(), attachment.getContentType(), IOUtils.toByteArray(attachment.getInputStream())));
+        for (DataSource attachmentDs : attachments) {
+            Attachment attachment = Attachment.fromDataSource(attachmentDs);
+            if (attachment != null) {
+                this.attachments.add(attachment);
             }
         }
     }