Skip to content

Commit cc8fb94

Browse files
authored
fix: Store UI reference on start (#22532) (#22538)
Store UI ref to have a correct ui for TransferProgress even when element is detached during transfer. Fixes #22243
1 parent 1f90fcd commit cc8fb94

File tree

13 files changed

+169
-23
lines changed

13 files changed

+169
-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
@@ -62,14 +62,15 @@ public class AbstractDownloadHandlerTest {
6262
private DownloadEvent downloadEvent;
6363
private ByteArrayOutputStream outputStream;
6464
private Element owner;
65+
private UI ui;
6566

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

72-
UI ui = Mockito.mock(UI.class);
73+
ui = Mockito.mock(UI.class);
7374
// run the command immediately
7475
Mockito.doAnswer(invocation -> {
7576
Command command = invocation.getArgument(0);
@@ -90,6 +91,7 @@ public void setUp() throws IOException {
9091
public void handleDownloadRequest(DownloadEvent event) {
9192
}
9293
};
94+
handler.setTransferUI(ui);
9395
mockContext = Mockito.mock(TransferContext.class);
9496
Mockito.when(mockContext.contentLength()).thenReturn(TOTAL_BYTES);
9597
listener = Mockito.mock(TransferProgressListener.class);
@@ -195,6 +197,7 @@ public void handleDownloadRequest(DownloadEvent event) {
195197
successAtomic.set(success);
196198
});
197199

200+
customHandler.setTransferUI(ui);
198201
customHandler.handleDownloadRequest(downloadEvent);
199202

200203
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
@@ -48,6 +48,7 @@ public class ClassDownloadHandlerTest {
4848
private DownloadEvent downloadEvent;
4949
private OutputStream outputStream;
5050
private Element owner;
51+
private UI ui;
5152

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

59-
UI ui = Mockito.mock(UI.class);
60+
ui = Mockito.mock(UI.class);
6061
// run the command immediately
6162
Mockito.doAnswer(invocation -> {
6263
Command command = invocation.getArgument(0);
@@ -137,6 +138,7 @@ public void transferProgressListener_addListener_errorOccured_errorlistenerInvok
137138
Mockito.when(event.getSession()).thenReturn(session);
138139
Mockito.when(event.getResponse()).thenReturn(response);
139140
Mockito.when(event.getOwningElement()).thenReturn(owner);
141+
Mockito.when(event.getUI()).thenReturn(ui);
140142
OutputStream outputStreamMock = Mockito.mock(OutputStream.class);
141143
Mockito.doThrow(new IOException("I/O exception")).when(outputStreamMock)
142144
.write(Mockito.any(byte[].class), Mockito.anyInt(),

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

Lines changed: 8 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

@@ -209,6 +212,7 @@ public void transferProgressListener_addListener_callbackUncheckedExceptionOccur
209212
Mockito.when(event.getSession()).thenReturn(session);
210213
Mockito.when(event.getResponse()).thenReturn(response);
211214
Mockito.when(event.getOwningElement()).thenReturn(owner);
215+
Mockito.when(event.getUI()).thenReturn(ui);
212216
OutputStream outputStreamMock = Mockito.mock(OutputStream.class);
213217
Mockito.when(event.getOutputStream()).thenReturn(outputStreamMock);
214218

@@ -241,6 +245,7 @@ public void transferProgressListener_addListener_callbackResponseError_errorList
241245
Mockito.when(event.getSession()).thenReturn(session);
242246
Mockito.when(event.getResponse()).thenReturn(response);
243247
Mockito.when(event.getOwningElement()).thenReturn(owner);
248+
Mockito.when(event.getUI()).thenReturn(ui);
244249
OutputStream outputStreamMock = Mockito.mock(OutputStream.class);
245250
Mockito.when(event.getOutputStream()).thenReturn(outputStreamMock);
246251

@@ -281,6 +286,7 @@ public void inline_setFileNameInvokedByDefault() throws IOException {
281286
Mockito.when(event.getResponse()).thenReturn(response);
282287
Mockito.when(event.getOwningElement()).thenReturn(owner);
283288
Mockito.when(event.getOutputStream()).thenReturn(outputStream);
289+
Mockito.when(event.getUI()).thenReturn(ui);
284290
Mockito.when(response.getOutputStream()).thenReturn(outputStream);
285291
Mockito.when(response.getService()).thenReturn(service);
286292
Mockito.when(service.getMimeType(Mockito.anyString()))
@@ -307,6 +313,7 @@ public void attachment_doesNotSetFileNameWhenInlined() throws IOException {
307313
Mockito.when(event.getResponse()).thenReturn(response);
308314
Mockito.when(event.getOwningElement()).thenReturn(owner);
309315
Mockito.when(event.getOutputStream()).thenReturn(outputStream);
316+
Mockito.when(event.getUI()).thenReturn(ui);
310317
Mockito.when(response.getOutputStream()).thenReturn(outputStream);
311318
Mockito.when(response.getService()).thenReturn(service);
312319
Mockito.when(service.getMimeType(Mockito.anyString()))

0 commit comments

Comments
 (0)