Skip to content

Commit 624d862

Browse files
vaadin-botcaalador
andauthored
fix: Store UI reference on start (#22532) (#22537)
Store UI ref to have a correct ui for TransferProgress even when element is detached during transfer. Fixes #22243 Co-authored-by: caalador <mikael.grankvist@vaadin.com>
1 parent 8705f81 commit 624d862

File tree

13 files changed

+170
-23
lines changed

13 files changed

+170
-23
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ public AbstractFileUploadHandler(FileUploadCallback successCallback,
5555
public void handleUploadRequest(UploadEvent event) throws IOException {
5656
UploadMetadata metadata = new UploadMetadata(event.getFileName(),
5757
event.getContentType(), event.getFileSize());
58+
setTransferUI(event.getUI());
5859
File file;
5960
try {
6061
file = fileFactory.createFile(metadata);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ public ClassDownloadHandler(Class<?> clazz, String resourceName,
9090
@Override
9191
public void handleDownloadRequest(DownloadEvent downloadEvent)
9292
throws IOException {
93+
setTransferUI(downloadEvent.getUI());
9394
if (clazz.getResource(resourceName) == null) {
9495
LoggerFactory.getLogger(ClassDownloadHandler.class)
9596
.warn("No resource found for '{}'", resourceName);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ public FileDownloadHandler(File file, String fileNameOverride) {
7373
@Override
7474
public void handleDownloadRequest(DownloadEvent downloadEvent)
7575
throws IOException {
76+
setTransferUI(downloadEvent.getUI());
7677
VaadinResponse response = downloadEvent.getResponse();
7778
try (OutputStream outputStream = downloadEvent.getOutputStream();
7879
FileInputStream inputStream = new FileInputStream(file)) {

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ public InMemoryUploadHandler(InMemoryUploadCallback successCallback) {
4040

4141
@Override
4242
public void handleUploadRequest(UploadEvent event) throws IOException {
43+
setTransferUI(event.getUI());
4344
byte[] data;
4445
try {
4546
try (InputStream inputStream = event.getInputStream();

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ public InputStreamDownloadHandler(InputStreamDownloadCallback callback) {
4949
public void handleDownloadRequest(DownloadEvent downloadEvent)
5050
throws IOException {
5151
VaadinResponse response = downloadEvent.getResponse();
52+
setTransferUI(downloadEvent.getUI());
5253
DownloadResponse download;
5354
try {
5455
download = callback.complete(downloadEvent);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@ public ServletResourceDownloadHandler(String path,
7676
@Override
7777
public void handleDownloadRequest(DownloadEvent downloadEvent)
7878
throws IOException {
79+
setTransferUI(downloadEvent.getUI());
7980
VaadinService service = downloadEvent.getRequest().getService();
8081
VaadinResponse response = downloadEvent.getResponse();
8182
if (service instanceof VaadinServletService servletService) {

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

Lines changed: 40 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import java.util.List;
2424
import java.util.Objects;
2525

26+
import com.vaadin.flow.component.UI;
2627
import com.vaadin.flow.function.SerializableBiConsumer;
2728
import com.vaadin.flow.function.SerializableConsumer;
2829
import com.vaadin.flow.function.SerializableRunnable;
@@ -55,6 +56,13 @@ public abstract class TransferProgressAwareHandler<T, R extends TransferProgress
5556
*/
5657
protected abstract TransferContext getTransferContext(T transferEvent);
5758

59+
/**
60+
* UI that the transfer was started with. Set from onStart(context).
61+
* <p>
62+
* protected for tests to set the ui.
63+
*/
64+
private UI ui;
65+
5866
/**
5967
* Adds a listener to be notified of data transfer progress events, such as:
6068
* <ul>
@@ -81,8 +89,9 @@ public abstract class TransferProgressAwareHandler<T, R extends TransferProgress
8189
protected Registration addTransferProgressListener(
8290
TransferProgressListener listener) {
8391
Objects.requireNonNull(listener, "Listener cannot be null");
84-
TransferProgressListener wrapper = new TransferProgressListenerWrapper(
92+
TransferProgressListenerWrapper wrapper = new TransferProgressListenerWrapper(
8593
listener);
94+
wrapper.setTransferUI(ui);
8695
return addTransferProgressListenerInternal(wrapper);
8796
}
8897

@@ -104,7 +113,7 @@ public R whenStart(SerializableRunnable startHandler) {
104113
addTransferProgressListenerInternal(new TransferProgressListener() {
105114
@Override
106115
public void onStart(TransferContext context) {
107-
context.getUI().access(startHandler::run);
116+
ui.access(startHandler::run);
108117
}
109118
});
110119
return (R) this;
@@ -129,7 +138,7 @@ public R whenStart(SerializableConsumer<TransferContext> startHandler) {
129138
addTransferProgressListenerInternal(new TransferProgressListener() {
130139
@Override
131140
public void onStart(TransferContext context) {
132-
context.getUI().access(() -> startHandler.accept(context));
141+
ui.access(() -> startHandler.accept(context));
133142
}
134143
});
135144
return (R) this;
@@ -161,7 +170,7 @@ public R onProgress(
161170
@Override
162171
public void onProgress(TransferContext context,
163172
long transferredBytes, long totalBytes) {
164-
context.getUI().access(() -> progressHandler.accept(context,
173+
ui.access(() -> progressHandler.accept(context,
165174
transferredBytes, totalBytes));
166175
}
167176

@@ -218,8 +227,8 @@ public R onProgress(SerializableBiConsumer<Long, Long> progressHandler,
218227
@Override
219228
public void onProgress(TransferContext context,
220229
long transferredBytes, long totalBytes) {
221-
context.getUI().access(() -> progressHandler
222-
.accept(transferredBytes, totalBytes));
230+
ui.access(() -> progressHandler.accept(transferredBytes,
231+
totalBytes));
223232
}
224233

225234
@Override
@@ -279,15 +288,13 @@ public R whenComplete(
279288
addTransferProgressListenerInternal(new TransferProgressListener() {
280289
@Override
281290
public void onError(TransferContext context, IOException reason) {
282-
context.getUI()
283-
.access(() -> completeOrTerminateHandler.accept(false));
291+
ui.access(() -> completeOrTerminateHandler.accept(false));
284292
}
285293

286294
@Override
287295
public void onComplete(TransferContext context,
288296
long transferredBytes) {
289-
context.getUI()
290-
.access(() -> completeOrTerminateHandler.accept(true));
297+
ui.access(() -> completeOrTerminateHandler.accept(true));
291298
}
292299
});
293300
return (R) this;
@@ -319,14 +326,14 @@ public R whenComplete(
319326
addTransferProgressListenerInternal(new TransferProgressListener() {
320327
@Override
321328
public void onError(TransferContext context, IOException reason) {
322-
context.getUI().access(() -> completeOrTerminateHandler
323-
.accept(context, false));
329+
ui.access(() -> completeOrTerminateHandler.accept(context,
330+
false));
324331
}
325332

326333
@Override
327334
public void onComplete(TransferContext context,
328335
long transferredBytes) {
329-
context.getUI().access(
336+
ui.access(
330337
() -> completeOrTerminateHandler.accept(context, true));
331338
}
332339
});
@@ -388,6 +395,16 @@ private Registration addTransferProgressListenerInternal(
388395
return Registration.addAndRemove(listeners, listener);
389396
}
390397

398+
protected void setTransferUI(UI ui) {
399+
this.ui = ui;
400+
if (listeners != null) {
401+
listeners.stream()
402+
.filter(l -> l instanceof TransferProgressListenerWrapper)
403+
.map(TransferProgressListenerWrapper.class::cast)
404+
.forEach(listener -> listener.setTransferUI(ui));
405+
}
406+
}
407+
391408
/**
392409
* A wrapper for {@link TransferProgressListener} that ensures that UI
393410
* updates in transfer progress listeners are pushed to client
@@ -396,6 +413,7 @@ private Registration addTransferProgressListenerInternal(
396413
private static final class TransferProgressListenerWrapper
397414
implements TransferProgressListener {
398415
private final TransferProgressListener delegate;
416+
private UI ui;
399417

400418
public TransferProgressListenerWrapper(
401419
TransferProgressListener delegate) {
@@ -404,30 +422,33 @@ public TransferProgressListenerWrapper(
404422

405423
@Override
406424
public void onStart(TransferContext context) {
407-
context.getUI().access(() -> delegate.onStart(context));
425+
ui.access(() -> delegate.onStart(context));
408426
}
409427

410428
@Override
411429
public void onProgress(TransferContext context, long transferredBytes,
412430
long totalBytes) {
413-
context.getUI().access(() -> delegate.onProgress(context,
414-
transferredBytes, totalBytes));
431+
ui.access(() -> delegate.onProgress(context, transferredBytes,
432+
totalBytes));
415433
}
416434

417435
@Override
418436
public void onError(TransferContext context, IOException reason) {
419-
context.getUI().access(() -> delegate.onError(context, reason));
437+
ui.access(() -> delegate.onError(context, reason));
420438
}
421439

422440
@Override
423441
public void onComplete(TransferContext context, long transferredBytes) {
424-
context.getUI().access(
425-
() -> delegate.onComplete(context, transferredBytes));
442+
ui.access(() -> delegate.onComplete(context, transferredBytes));
426443
}
427444

428445
@Override
429446
public long progressReportInterval() {
430447
return delegate.progressReportInterval();
431448
}
449+
450+
protected void setTransferUI(UI ui) {
451+
this.ui = ui;
452+
}
432453
}
433454
}

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,14 +63,15 @@ public class AbstractDownloadHandlerTest {
6363
private DownloadEvent downloadEvent;
6464
private ByteArrayOutputStream outputStream;
6565
private Element owner;
66+
private UI ui;
6667

6768
@Before
6869
public void setUp() throws IOException {
6970
request = Mockito.mock(VaadinRequest.class);
7071
response = Mockito.mock(VaadinResponse.class);
7172
session = Mockito.mock(VaadinSession.class);
7273

73-
UI ui = Mockito.mock(UI.class);
74+
ui = Mockito.mock(UI.class);
7475
// run the command immediately
7576
Mockito.doAnswer(invocation -> {
7677
Command command = invocation.getArgument(0);
@@ -91,6 +92,7 @@ public void setUp() throws IOException {
9192
public void handleDownloadRequest(DownloadEvent event) {
9293
}
9394
};
95+
handler.setTransferUI(ui);
9496
mockContext = Mockito.mock(TransferContext.class);
9597
Mockito.when(mockContext.contentLength()).thenReturn(TOTAL_BYTES);
9698
listener = Mockito.mock(TransferProgressListener.class);
@@ -196,6 +198,7 @@ public void handleDownloadRequest(DownloadEvent event) {
196198
successAtomic.set(success);
197199
});
198200

201+
customHandler.setTransferUI(ui);
199202
customHandler.handleDownloadRequest(downloadEvent);
200203

201204
Assert.assertTrue(successAtomic.get());

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ public class ClassDownloadHandlerTest {
4949
private DownloadEvent downloadEvent;
5050
private OutputStream outputStream;
5151
private Element owner;
52+
private UI ui;
5253

5354
@Before
5455
public void setUp() throws IOException {
@@ -57,7 +58,7 @@ public void setUp() throws IOException {
5758
session = Mockito.mock(VaadinSession.class);
5859
service = Mockito.mock(VaadinService.class);
5960

60-
UI ui = Mockito.mock(UI.class);
61+
ui = Mockito.mock(UI.class);
6162
// run the command immediately
6263
Mockito.doAnswer(invocation -> {
6364
Command command = invocation.getArgument(0);
@@ -139,6 +140,7 @@ public void transferProgressListener_addListener_errorOccured_errorlistenerInvok
139140
Mockito.when(event.getSession()).thenReturn(session);
140141
Mockito.when(event.getResponse()).thenReturn(response);
141142
Mockito.when(event.getOwningElement()).thenReturn(owner);
143+
Mockito.when(event.getUI()).thenReturn(ui);
142144
OutputStream outputStreamMock = Mockito.mock(OutputStream.class);
143145
Mockito.doThrow(new IOException("I/O exception")).when(outputStreamMock)
144146
.write(Mockito.any(byte[].class), Mockito.anyInt(),

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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ public class InputStreamDownloadHandlerTest {
4949
private DownloadEvent downloadEvent;
5050
private OutputStream outputStream;
5151
private Element owner;
52+
private UI ui;
5253

5354
@Before
5455
public void setUp() throws IOException {
@@ -57,7 +58,7 @@ public void setUp() throws IOException {
5758
session = Mockito.mock(VaadinSession.class);
5859
service = Mockito.mock(VaadinService.class);
5960

60-
UI ui = Mockito.mock(UI.class);
61+
ui = Mockito.mock(UI.class);
6162
// run the command immediately
6263
Mockito.doAnswer(invocation -> {
6364
Command command = invocation.getArgument(0);
@@ -140,6 +141,7 @@ public void transferProgressListener_addListener_errorOccurred_errorListenerInvo
140141
Mockito.when(event.getSession()).thenReturn(session);
141142
Mockito.when(event.getResponse()).thenReturn(response);
142143
Mockito.when(event.getOwningElement()).thenReturn(owner);
144+
Mockito.when(event.getUI()).thenReturn(ui);
143145
OutputStream outputStreamMock = Mockito.mock(OutputStream.class);
144146
Mockito.doThrow(new IOException("I/O exception")).when(outputStreamMock)
145147
.write(Mockito.any(byte[].class), Mockito.anyInt(),
@@ -178,6 +180,7 @@ public void transferProgressListener_addListener_callbackIOExceptionOccurred_err
178180
Mockito.when(event.getSession()).thenReturn(session);
179181
Mockito.when(event.getResponse()).thenReturn(response);
180182
Mockito.when(event.getOwningElement()).thenReturn(owner);
183+
Mockito.when(event.getUI()).thenReturn(ui);
181184
OutputStream outputStreamMock = Mockito.mock(OutputStream.class);
182185
Mockito.when(event.getOutputStream()).thenReturn(outputStreamMock);
183186

@@ -210,6 +213,7 @@ public void transferProgressListener_addListener_callbackUncheckedExceptionOccur
210213
Mockito.when(event.getSession()).thenReturn(session);
211214
Mockito.when(event.getResponse()).thenReturn(response);
212215
Mockito.when(event.getOwningElement()).thenReturn(owner);
216+
Mockito.when(event.getUI()).thenReturn(ui);
213217
OutputStream outputStreamMock = Mockito.mock(OutputStream.class);
214218
Mockito.when(event.getOutputStream()).thenReturn(outputStreamMock);
215219

@@ -243,6 +247,7 @@ public void transferProgressListener_addListener_callbackResponseError_errorList
243247
Mockito.when(event.getSession()).thenReturn(session);
244248
Mockito.when(event.getResponse()).thenReturn(response);
245249
Mockito.when(event.getOwningElement()).thenReturn(owner);
250+
Mockito.when(event.getUI()).thenReturn(ui);
246251
OutputStream outputStreamMock = Mockito.mock(OutputStream.class);
247252
Mockito.when(event.getOutputStream()).thenReturn(outputStreamMock);
248253

@@ -272,6 +277,7 @@ public void transferProgressListener_addListener_callbackResponseException_error
272277
Mockito.when(event.getSession()).thenReturn(session);
273278
Mockito.when(event.getResponse()).thenReturn(response);
274279
Mockito.when(event.getOwningElement()).thenReturn(owner);
280+
Mockito.when(event.getUI()).thenReturn(ui);
275281
OutputStream outputStreamMock = Mockito.mock(OutputStream.class);
276282
Mockito.when(event.getOutputStream()).thenReturn(outputStreamMock);
277283

@@ -323,6 +329,7 @@ public void inline_setFileNameInvokedByDefault() throws IOException {
323329
Mockito.when(event.getResponse()).thenReturn(response);
324330
Mockito.when(event.getOwningElement()).thenReturn(owner);
325331
Mockito.when(event.getOutputStream()).thenReturn(outputStream);
332+
Mockito.when(event.getUI()).thenReturn(ui);
326333
Mockito.when(response.getOutputStream()).thenReturn(outputStream);
327334
Mockito.when(response.getService()).thenReturn(service);
328335
Mockito.when(service.getMimeType(Mockito.anyString()))
@@ -349,6 +356,7 @@ public void attachment_doesNotSetFileNameWhenInlined() throws IOException {
349356
Mockito.when(event.getResponse()).thenReturn(response);
350357
Mockito.when(event.getOwningElement()).thenReturn(owner);
351358
Mockito.when(event.getOutputStream()).thenReturn(outputStream);
359+
Mockito.when(event.getUI()).thenReturn(ui);
352360
Mockito.when(response.getOutputStream()).thenReturn(outputStream);
353361
Mockito.when(response.getService()).thenReturn(service);
354362
Mockito.when(service.getMimeType(Mockito.anyString()))

0 commit comments

Comments
 (0)