provide the file name to the AttachmentHandler also if there is one i…

…n the content-disposition header (#789)
HtmlUnit · May 16, 2024 · ac60725 · ac60725
1 parent 13beed0
commit ac60725
Show file tree

Hide file tree

Showing 3 changed files with 93 additions and 8 deletions.
diff --git a/src/main/java/org/htmlunit/WebClient.java b/src/main/java/org/htmlunit/WebClient.java
@@ -59,6 +59,7 @@
 import org.apache.http.NoHttpResponseException;
 import org.apache.http.client.CredentialsProvider;
 import org.apache.http.cookie.MalformedCookieException;
+import org.htmlunit.attachment.Attachment;
 import org.htmlunit.attachment.AttachmentHandler;
 import org.htmlunit.csp.Policy;
 import org.htmlunit.csp.url.URI;
@@ -565,7 +566,7 @@ public Page loadWebResponseInto(final WebResponse webResponse, final WebWindow w
      * @see #setAttachmentHandler(AttachmentHandler)
      */
     public Page loadWebResponseInto(final WebResponse webResponse, final WebWindow webWindow,
-            final String forceAttachmentWithFilename)
+            String forceAttachmentWithFilename)
             throws IOException, FailingHttpStatusCodeException {
         WebAssert.notNull("webResponse", webResponse);
         WebAssert.notNull("webWindow", webWindow);
@@ -581,6 +582,13 @@ public Page loadWebResponseInto(final WebResponse webResponse, final WebWindow w
 
         if (attachmentHandler_ != null
                 && (forceAttachmentWithFilename != null || attachmentHandler_.isAttachment(webResponse))) {
+
+            // check content disposition header for nothing provided
+            if (StringUtils.isEmpty(forceAttachmentWithFilename)) {
+                final String disp = webResponse.getResponseHeaderValue(HttpHeader.CONTENT_DISPOSITION);
+                forceAttachmentWithFilename = Attachment.getSuggestedFilename(disp);
+            }
+
             if (attachmentHandler_.handleAttachment(webResponse,
                         StringUtils.isEmpty(forceAttachmentWithFilename) ? null : forceAttachmentWithFilename)) {
                 // the handling is done by the attachment handler;

diff --git a/src/main/java/org/htmlunit/attachment/Attachment.java b/src/main/java/org/htmlunit/attachment/Attachment.java
@@ -74,27 +74,38 @@ public String getSuggestedFilename() {
 
         final WebResponse response = page_.getWebResponse();
         final String disp = response.getResponseHeaderValue(HttpHeader.CONTENT_DISPOSITION);
-        if (disp == null) {
+        return getSuggestedFilename(disp);
+    }
+
+    /**
+     * Returns the attachment's filename, as suggested by the <code>Content-Disposition</code>
+     * header, or {@code null} if no filename was suggested.
+     * @param contentDispositonHeader the <code>Content-Disposition</code> header
+     *
+     * @return the attachment's suggested filename, or {@code null} if none was suggested
+     */
+    public static String getSuggestedFilename(final String contentDispositonHeader) {
+        if (contentDispositonHeader == null) {
             return null;
         }
 
-        int start = disp.indexOf("filename=");
+        int start = contentDispositonHeader.indexOf("filename=");
         if (start == -1) {
             return null;
         }
         start += "filename=".length();
-        if (start >= disp.length()) {
+        if (start >= contentDispositonHeader.length()) {
             return null;
         }
 
-        int end = disp.indexOf(';', start);
+        int end = contentDispositonHeader.indexOf(';', start);
         if (end == -1) {
-            end = disp.length();
+            end = contentDispositonHeader.length();
         }
-        if (disp.charAt(start) == '"' && disp.charAt(end - 1) == '"') {
+        if (contentDispositonHeader.charAt(start) == '"' && contentDispositonHeader.charAt(end - 1) == '"') {
             start++;
             end--;
         }
-        return disp.substring(start, end);
+        return contentDispositonHeader.substring(start, end);
     }
 }
diff --git a/src/test/java/org/htmlunit/attachment/AttachmentTest.java b/src/test/java/org/htmlunit/attachment/AttachmentTest.java
@@ -325,11 +325,13 @@ public void handleResponseFromHandler() throws Exception {
 
         final WebClient client = getWebClient();
         final List<WebResponse> attachments = new ArrayList<>();
+        final List<String> attachmentFilenames = new ArrayList<>();
 
         client.setAttachmentHandler(new AttachmentHandler() {
             @Override
             public boolean handleAttachment(final WebResponse response, final String attachmentFilename) {
                 attachments.add(response);
+                attachmentFilenames.add(attachmentFilename);
                 return true;
             }
 
@@ -347,12 +349,14 @@ public void handleAttachment(final Page page, final String attachmentFilename) {
         conn.setResponse(URL_SECOND, content2, 200, "OK", MimeType.TEXT_HTML, headers);
         client.setWebConnection(conn);
         assertTrue(attachments.isEmpty());
+        assertTrue(attachmentFilenames.isEmpty());
 
         final HtmlPage result = client.getPage(URL_FIRST);
         final HtmlAnchor anchor = result.getAnchors().get(0);
         final Page clickResult = anchor.click();
         assertEquals(result, clickResult);
         assertEquals(1, attachments.size());
+        assertEquals(1, attachmentFilenames.size());
         assertEquals(1, client.getWebWindows().size());
 
         final WebResponse attachmentResponse = attachments.get(0);
@@ -361,5 +365,67 @@ public void handleAttachment(final Page page, final String attachmentFilename) {
         assertEquals(MimeType.TEXT_HTML, attachmentResponse.getContentType());
         assertEquals(200, attachmentResponse.getStatusCode());
         assertEquals(URL_SECOND, attachmentResponse.getWebRequest().getUrl());
+
+        assertNull(attachmentFilenames.get(0));
+    }
+
+    /**
+     * Tests attachment callbacks and the contents of attachments.
+     * @throws Exception if an error occurs
+     */
+    @Test
+    public void handleResponseFromHandlerWithFileName() throws Exception {
+        final String content1 = "<html><body>\n"
+            + "<form method='POST' name='form' action='" + URL_SECOND + "'>\n"
+            + "<input type='submit' value='ok'>\n"
+            + "</form>\n"
+            + "<a href='#' onclick='document.form.submit()'>click me</a>\n"
+            + "</body></html>";
+        final String content2 = "download file contents";
+
+        final WebClient client = getWebClient();
+        final List<WebResponse> attachments = new ArrayList<>();
+        final List<String> attachmentFilenames = new ArrayList<>();
+
+        client.setAttachmentHandler(new AttachmentHandler() {
+            @Override
+            public boolean handleAttachment(final WebResponse response, final String attachmentFilename) {
+                attachments.add(response);
+                attachmentFilenames.add(attachmentFilename);
+                return true;
+            }
+
+            @Override
+            public void handleAttachment(final Page page, final String attachmentFilename) {
+                throw new IllegalAccessError("handleAttachment(Page, String) called");
+            }
+        });
+
+        final List<NameValuePair> headers = new ArrayList<>();
+        headers.add(new NameValuePair("Content-Disposition", "attachment; filename=htmlunit.zip"));
+
+        final MockWebConnection conn = new MockWebConnection();
+        conn.setResponse(URL_FIRST, content1);
+        conn.setResponse(URL_SECOND, content2, 200, "OK", MimeType.TEXT_HTML, headers);
+        client.setWebConnection(conn);
+        assertTrue(attachments.isEmpty());
+        assertTrue(attachmentFilenames.isEmpty());
+
+        final HtmlPage result = client.getPage(URL_FIRST);
+        final HtmlAnchor anchor = result.getAnchors().get(0);
+        final Page clickResult = anchor.click();
+        assertEquals(result, clickResult);
+        assertEquals(1, attachments.size());
+        assertEquals(1, attachmentFilenames.size());
+        assertEquals(1, client.getWebWindows().size());
+
+        final WebResponse attachmentResponse = attachments.get(0);
+        final InputStream attachmentStream = attachmentResponse.getContentAsStream();
+        HttpWebConnectionTest.assertEquals(new ByteArrayInputStream(content2.getBytes()), attachmentStream);
+        assertEquals(MimeType.TEXT_HTML, attachmentResponse.getContentType());
+        assertEquals(200, attachmentResponse.getStatusCode());
+        assertEquals(URL_SECOND, attachmentResponse.getWebRequest().getUrl());
+
+        assertEquals("htmlunit.zip", attachmentFilenames.get(0));
     }
 }