JAMES-1975 Refactor DKIM

apache · Jun 18, 2019 · b8f9199 · b8f9199
1 parent 30a07f2
commit b8f9199
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 49 deletions.
diff --git a/server/mailet/dkim/src/main/java/org/apache/james/jdkim/mailets/DKIMSign.java b/server/mailet/dkim/src/main/java/org/apache/james/jdkim/mailets/DKIMSign.java
@@ -32,14 +32,15 @@
 import java.security.Security;
 import java.security.spec.InvalidKeySpecException;
 import java.security.spec.RSAPrivateCrtKeySpec;
-import java.util.Enumeration;
-import java.util.LinkedList;
+import java.util.Collections;
 import java.util.List;
 
 import javax.mail.Header;
 import javax.mail.MessagingException;
 import javax.mail.internet.MimeMessage;
 
+import com.github.fge.lambdas.Throwing;
+import com.google.common.collect.ImmutableList;
 import org.apache.commons.io.IOUtils;
 import org.apache.james.jdkim.DKIMSigner;
 import org.apache.james.jdkim.api.BodyHasher;
@@ -142,43 +143,37 @@ public void service(Mail mail) throws MessagingException {
                     os = new CRLFOutputStream(os);
                 }
                 message.writeTo(os);
-                bhj.getOutputStream().close();
             } catch (IOException e) {
-                throw new MessagingException("Exception calculating bodyhash: "
-                        + e.getMessage(), e);
+                throw new MessagingException("Exception calculating bodyhash: " + e.getMessage(), e);
+            } finally {
+                try {
+                    bhj.getOutputStream().close();
+                } catch (IOException e) {
+                    throw new MessagingException("Exception calculating bodyhash: " + e.getMessage(), e);
+                }
             }
             String signatureHeader = signer.sign(headers, bhj);
             // Unfortunately JavaMail does not give us a method to add headers
             // on top.
             // message.addHeaderLine(signatureHeader);
             prependHeader(message, signatureHeader);
         } catch (PermFailException e) {
-            throw new MessagingException("PermFail while signing: "
-                    + e.getMessage(), e);
+            throw new MessagingException("PermFail while signing: " + e.getMessage(), e);
         }
 
     }
 
     private void prependHeader(MimeMessage message, String signatureHeader)
             throws MessagingException {
-        List<String> prevHeader = new LinkedList<>();
-        // read all the headers
-        for (Enumeration<String> e = message.getAllHeaderLines(); e.hasMoreElements();) {
-            String headerLine = e.nextElement();
-            prevHeader.add(headerLine);
-        }
-        // remove all the headers
-        for (Enumeration<Header> e = message.getAllHeaders(); e.hasMoreElements();) {
-            Header header = e.nextElement();
-            message.removeHeader(header.getName());
-        }
-        // add our header
+        List<String> prevHeader = Collections.list(message.getAllHeaderLines());
+        Collections.list(message.getAllHeaders())
+            .stream()
+            .map(Header::getName)
+            .forEach(Throwing.consumer(message::removeHeader).sneakyThrow());
+
         message.addHeaderLine(signatureHeader);
-        // add the remaining headers using "addHeaderLine" that won't alter the
-        // insertion order.
-        for (String header : prevHeader) {
-            message.addHeaderLine(header);
-        }
+        prevHeader
+            .forEach(Throwing.consumer(message::addHeaderLine).sneakyThrow());
     }
 
     private PrivateKey extractPrivateKey(InputStream rawKey, char[] passphrase) throws IOException, NoSuchAlgorithmException, InvalidKeySpecException {

diff --git a/server/mailet/dkim/src/main/java/org/apache/james/jdkim/mailets/DKIMVerify.java b/server/mailet/dkim/src/main/java/org/apache/james/jdkim/mailets/DKIMVerify.java
@@ -106,12 +106,20 @@ static List<SignatureRecord> verify(DKIMVerifier verifier, MimeMessage message,
                     os = new CRLFOutputStream(os);
                 }
                 message.writeTo(os);
-                bh.getOutputStream().close();
             }
 
         } catch (IOException e) {
             throw new MessagingException("Exception calculating bodyhash: "
                     + e.getMessage(), e);
+        } finally {
+            try {
+                if (bh != null) {
+                    bh.getOutputStream().close();
+                }
+            } catch (IOException e) {
+                throw new MessagingException("Exception calculating bodyhash: "
+                        + e.getMessage(), e);
+            }
         }
         return verifier.verify(bh);
     }

diff --git a/server/mailet/dkim/src/main/java/org/apache/james/jdkim/mailets/MimeMessageHeaders.java b/server/mailet/dkim/src/main/java/org/apache/james/jdkim/mailets/MimeMessageHeaders.java
@@ -19,51 +19,60 @@
 
 package org.apache.james.jdkim.mailets;
 
-import java.util.Enumeration;
-import java.util.HashMap;
-import java.util.LinkedList;
 import java.util.List;
-import java.util.Map;
+import java.util.Locale;
 
 import javax.mail.MessagingException;
 import javax.mail.internet.MimeMessage;
 
+import org.apache.commons.lang3.tuple.Pair;
 import org.apache.james.jdkim.api.Headers;
 
+import com.github.fge.lambdas.Throwing;
+import com.github.steveash.guavate.Guavate;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableListMultimap;
+import com.google.common.collect.Iterators;
+import com.google.common.collect.Streams;
+
 /**
  * An adapter to let DKIMSigner read headers from MimeMessage
  */
 final class MimeMessageHeaders implements Headers {
 
-    private final Map<String, List<String>> headers;
+    private final ImmutableListMultimap<String, String> headers;
     private final List<String> fields;
 
     public MimeMessageHeaders(MimeMessage message) throws MessagingException {
-        headers = new HashMap<>();
-        fields = new LinkedList<>();
-        for (Enumeration<String> e = message.getAllHeaderLines(); e.hasMoreElements();) {
-            String head = e.nextElement();
-            int p = head.indexOf(':');
-            if (p <= 0) {
-                throw new MessagingException("Bad header line: " + head);
-            }
-            String headerName = head.substring(0, p).trim();
-            String headerNameLC = headerName.toLowerCase();
-            fields.add(headerName);
-            List<String> strings = headers.get(headerNameLC);
-            if (strings == null) {
-                strings = new LinkedList<>();
-                headers.put(headerNameLC, strings);
-            }
-            strings.add(head);
-        }
+        ImmutableList<Pair<String, String>> headsAndLines = Streams.stream(Iterators.forEnumeration(message.getAllHeaderLines()))
+                .map(Throwing.function(this::extractHeaderLine).sneakyThrow())
+                .collect(Guavate.toImmutableList());
+
+        fields = headsAndLines
+            .stream()
+            .map(Pair::getKey)
+            .collect(Guavate.toImmutableList());
+
+        headers = headsAndLines
+            .stream()
+            .collect(Guavate.toImmutableListMultimap(
+                pair -> pair.getKey().toLowerCase(Locale.US),
+                Pair::getValue));
     }
 
     public List<String> getFields() {
         return fields;
     }
 
     public List<String> getFields(String name) {
-        return headers.get(name.toLowerCase());
+        return headers.get(name.toLowerCase(Locale.US));
+    }
+
+    private Pair<String, String> extractHeaderLine(String header) throws MessagingException {
+        int fieldSeperatorPosition = header.indexOf(':');
+        if (fieldSeperatorPosition <= 0) {
+            throw new MessagingException("Bad header line: " + header);
+        }
+        return Pair.of(header.substring(0, fieldSeperatorPosition).trim(), header);
     }
 }