fix: notify listener for all InputStreamDownloadCallback failures (#21994) (#21995)

Makes sure that transfer progress listeners are notified also
for DownloadResponse.error return value and unchecked exception
thrown by InputStreamDownloadCallback.

Fixes #21931

Co-authored-by: Marco Collovati <[email protected]>

Author: vaadin-bot

flow-server/src/main/java/com/vaadin/flow/server/streams/InputStreamDownloadHandler.java

@@ -19,6 +19,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.io.UncheckedIOException;
 
 import com.vaadin.flow.server.HttpStatusCode;
 import com.vaadin.flow.server.VaadinResponse;
@@ -51,14 +52,27 @@ public void handleDownloadRequest(DownloadEvent downloadEvent)
         DownloadResponse download;
         try {
             download = callback.complete(downloadEvent);
-        } catch (IOException ioe) {
+        } catch (IOException | RuntimeException e) {
             // Set status before output is closed (see #8740)
             response.setStatus(HttpStatusCode.INTERNAL_SERVER_ERROR.getCode());
-            notifyError(downloadEvent, ioe);
-            throw ioe;
+            IOException cause;
+            if (e instanceof IOException ioe) {
+                cause = ioe;
+            } else if (e instanceof UncheckedIOException uioe) {
+                cause = uioe.getCause();
+            } else {
+                cause = new IOException(e.getMessage(), e);
+            }
+            notifyError(downloadEvent, cause);
+            throw e;
         }
         if (download.hasError()) {
             response.setStatus(download.getError());
+            String message = download.getErrorMessage();
+            if (message == null) {
+                message = "Download failed with code " + download.getError();
+            }
+            notifyError(downloadEvent, new IOException(message));
             return;
         }
 

flow-server/src/test/java/com/vaadin/flow/server/streams/InputStreamDownloadHandlerTest.java

@@ -134,7 +134,7 @@ public void onError(TransferContext context, IOException reason) {
     }
 
     @Test
-    public void transferProgressListener_addListener_errorOccured_errorlistenerInvoked()
+    public void transferProgressListener_addListener_errorOccurred_errorListenerInvoked()
             throws IOException {
         DownloadEvent event = Mockito.mock(DownloadEvent.class);
         Mockito.when(event.getSession()).thenReturn(session);
@@ -173,7 +173,7 @@ public void transferProgressListener_addListener_errorOccured_errorlistenerInvok
     }
 
     @Test
-    public void transferProgressListener_addListener_callbackErrorOccured_errorlistenerInvoked() {
+    public void transferProgressListener_addListener_callbackIOExceptionOccurred_errorListenerInvoked() {
         DownloadEvent event = Mockito.mock(DownloadEvent.class);
         Mockito.when(event.getSession()).thenReturn(session);
         Mockito.when(event.getResponse()).thenReturn(response);
@@ -203,6 +203,65 @@ public void transferProgressListener_addListener_callbackErrorOccured_errorliste
                 whenCompleteResult.get());
     }
 
+    @Test
+    public void transferProgressListener_addListener_callbackUncheckedExceptionOccurred_errorListenerInvoked() {
+        DownloadEvent event = Mockito.mock(DownloadEvent.class);
+        Mockito.when(event.getSession()).thenReturn(session);
+        Mockito.when(event.getResponse()).thenReturn(response);
+        Mockito.when(event.getOwningElement()).thenReturn(owner);
+        OutputStream outputStreamMock = Mockito.mock(OutputStream.class);
+        Mockito.when(event.getOutputStream()).thenReturn(outputStreamMock);
+
+        AtomicReference<Boolean> whenCompleteResult = new AtomicReference<>();
+
+        InvocationTrackingTransferProgressListener transferListener = new InvocationTrackingTransferProgressListener();
+        DownloadHandler handler = DownloadHandler.fromInputStream(req -> {
+            throw new RuntimeException("I/O exception");
+        }, transferListener).whenComplete((context, success) -> {
+            whenCompleteResult.set(success);
+        });
+
+        try {
+            handler.handleDownloadRequest(event);
+            Assert.fail("Expected an exception to be thrown");
+        } catch (Exception e) {
+        }
+        Assert.assertEquals(List.of("onError"), transferListener.invocations);
+        Assert.assertNotNull("Expected whenComplete to be invoked, but was not",
+                whenCompleteResult.get());
+        Assert.assertFalse(
+                "Expected whenComplete to be invoked with false result, but got true",
+                whenCompleteResult.get());
+    }
+
+    @Test
+    public void transferProgressListener_addListener_callbackResponseError_errorListenerInvoked()
+            throws IOException {
+        DownloadEvent event = Mockito.mock(DownloadEvent.class);
+        Mockito.when(event.getSession()).thenReturn(session);
+        Mockito.when(event.getResponse()).thenReturn(response);
+        Mockito.when(event.getOwningElement()).thenReturn(owner);
+        OutputStream outputStreamMock = Mockito.mock(OutputStream.class);
+        Mockito.when(event.getOutputStream()).thenReturn(outputStreamMock);
+
+        AtomicReference<Boolean> whenCompleteResult = new AtomicReference<>();
+
+        InvocationTrackingTransferProgressListener transferListener = new InvocationTrackingTransferProgressListener();
+        DownloadHandler handler = DownloadHandler.fromInputStream(
+                req -> DownloadResponse.error(500, "I/O exception"),
+                transferListener).whenComplete((context, success) -> {
+                    whenCompleteResult.set(success);
+                });
+
+        handler.handleDownloadRequest(event);
+        Assert.assertEquals(List.of("onError"), transferListener.invocations);
+        Assert.assertNotNull("Expected whenComplete to be invoked, but was not",
+                whenCompleteResult.get());
+        Assert.assertFalse(
+                "Expected whenComplete to be invoked with false result, but got true",
+                whenCompleteResult.get());
+    }
+
     @Test
     public void inline_setFileNameInvokedByDefault() throws IOException {
         DownloadEvent event = Mockito.mock(DownloadEvent.class);