Skip to content

Commit 4ba18ff

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

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
@@ -54,6 +54,7 @@ public AbstractFileUploadHandler(FileUploadCallback successCallback,
5454
public void handleUploadRequest(UploadEvent event) throws IOException {
5555
UploadMetadata metadata = new UploadMetadata(event.getFileName(),
5656
event.getContentType(), event.getFileSize());
57+
setTransferUI(event.getUI());
5758
File file;
5859
try {
5960
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
@@ -87,6 +87,7 @@ public ClassDownloadHandler(Class<?> clazz, String resourceName,
8787
@Override
8888
public void handleDownloadRequest(DownloadEvent downloadEvent)
8989
throws IOException {
90+
setTransferUI(downloadEvent.getUI());
9091
if (clazz.getResource(resourceName) == null) {
9192
LoggerFactory.getLogger(ClassDownloadHandler.class)
9293
.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
@@ -72,6 +72,7 @@ public FileDownloadHandler(File file, String fileNameOverride) {
7272
@Override
7373
public void handleDownloadRequest(DownloadEvent downloadEvent)
7474
throws IOException {
75+
setTransferUI(downloadEvent.getUI());
7576
VaadinResponse response = downloadEvent.getResponse();
7677
try (OutputStream outputStream = downloadEvent.getOutputStream();
7778
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
@@ -37,6 +37,7 @@ public InMemoryUploadHandler(InMemoryUploadCallback successCallback) {
3737

3838
@Override
3939
public void handleUploadRequest(UploadEvent event) throws IOException {
40+
setTransferUI(event.getUI());
4041
byte[] data;
4142
try {
4243
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
@@ -48,6 +48,7 @@ public InputStreamDownloadHandler(InputStreamDownloadCallback callback) {
4848
public void handleDownloadRequest(DownloadEvent downloadEvent)
4949
throws IOException {
5050
VaadinResponse response = downloadEvent.getResponse();
51+
setTransferUI(downloadEvent.getUI());
5152
DownloadResponse download;
5253
try {
5354
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
@@ -75,6 +75,7 @@ public ServletResourceDownloadHandler(String path,
7575
@Override
7676
public void handleDownloadRequest(DownloadEvent downloadEvent)
7777
throws IOException {
78+
setTransferUI(downloadEvent.getUI());
7879
VaadinService service = downloadEvent.getRequest().getService();
7980
VaadinResponse response = downloadEvent.getResponse();
8081
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
@@ -22,6 +22,7 @@
2222
import java.util.List;
2323
import java.util.Objects;
2424

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

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

@@ -103,7 +112,7 @@ public R whenStart(SerializableRunnable startHandler) {
103112
addTransferProgressListenerInternal(new TransferProgressListener() {
104113
@Override
105114
public void onStart(TransferContext context) {
106-
context.getUI().access(startHandler::run);
115+
ui.access(startHandler::run);
107116
}
108117
});
109118
return (R) this;
@@ -128,7 +137,7 @@ public R whenStart(SerializableConsumer<TransferContext> startHandler) {
128137
addTransferProgressListenerInternal(new TransferProgressListener() {
129138
@Override
130139
public void onStart(TransferContext context) {
131-
context.getUI().access(() -> startHandler.accept(context));
140+
ui.access(() -> startHandler.accept(context));
132141
}
133142
});
134143
return (R) this;
@@ -160,7 +169,7 @@ public R onProgress(
160169
@Override
161170
public void onProgress(TransferContext context,
162171
long transferredBytes, long totalBytes) {
163-
context.getUI().access(() -> progressHandler.accept(context,
172+
ui.access(() -> progressHandler.accept(context,
164173
transferredBytes, totalBytes));
165174
}
166175

@@ -217,8 +226,8 @@ public R onProgress(SerializableBiConsumer<Long, Long> progressHandler,
217226
@Override
218227
public void onProgress(TransferContext context,
219228
long transferredBytes, long totalBytes) {
220-
context.getUI().access(() -> progressHandler
221-
.accept(transferredBytes, totalBytes));
229+
ui.access(() -> progressHandler.accept(transferredBytes,
230+
totalBytes));
222231
}
223232

224233
@Override
@@ -278,15 +287,13 @@ public R whenComplete(
278287
addTransferProgressListenerInternal(new TransferProgressListener() {
279288
@Override
280289
public void onError(TransferContext context, IOException reason) {
281-
context.getUI()
282-
.access(() -> completeOrTerminateHandler.accept(false));
290+
ui.access(() -> completeOrTerminateHandler.accept(false));
283291
}
284292

285293
@Override
286294
public void onComplete(TransferContext context,
287295
long transferredBytes) {
288-
context.getUI()
289-
.access(() -> completeOrTerminateHandler.accept(true));
296+
ui.access(() -> completeOrTerminateHandler.accept(true));
290297
}
291298
});
292299
return (R) this;
@@ -318,14 +325,14 @@ public R whenComplete(
318325
addTransferProgressListenerInternal(new TransferProgressListener() {
319326
@Override
320327
public void onError(TransferContext context, IOException reason) {
321-
context.getUI().access(() -> completeOrTerminateHandler
322-
.accept(context, false));
328+
ui.access(() -> completeOrTerminateHandler.accept(context,
329+
false));
323330
}
324331

325332
@Override
326333
public void onComplete(TransferContext context,
327334
long transferredBytes) {
328-
context.getUI().access(
335+
ui.access(
329336
() -> completeOrTerminateHandler.accept(context, true));
330337
}
331338
});
@@ -387,6 +394,16 @@ private Registration addTransferProgressListenerInternal(
387394
return Registration.addAndRemove(listeners, listener);
388395
}
389396

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

399417
public TransferProgressListenerWrapper(
400418
TransferProgressListener delegate) {
@@ -403,30 +421,33 @@ public TransferProgressListenerWrapper(
403421

404422
@Override
405423
public void onStart(TransferContext context) {
406-
context.getUI().access(() -> delegate.onStart(context));
424+
ui.access(() -> delegate.onStart(context));
407425
}
408426

409427
@Override
410428
public void onProgress(TransferContext context, long transferredBytes,
411429
long totalBytes) {
412-
context.getUI().access(() -> delegate.onProgress(context,
413-
transferredBytes, totalBytes));
430+
ui.access(() -> delegate.onProgress(context, transferredBytes,
431+
totalBytes));
414432
}
415433

416434
@Override
417435
public void onError(TransferContext context, IOException reason) {
418-
context.getUI().access(() -> delegate.onError(context, reason));
436+
ui.access(() -> delegate.onError(context, reason));
419437
}
420438

421439
@Override
422440
public void onComplete(TransferContext context, long transferredBytes) {
423-
context.getUI().access(
424-
() -> delegate.onComplete(context, transferredBytes));
441+
ui.access(() -> delegate.onComplete(context, transferredBytes));
425442
}
426443

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

‎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
@@ -60,14 +60,15 @@ public class AbstractDownloadHandlerTest {
6060
private DownloadEvent downloadEvent;
6161
private ByteArrayOutputStream outputStream;
6262
private Element owner;
63+
private UI ui;
6364

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

70-
UI ui = Mockito.mock(UI.class);
71+
ui = Mockito.mock(UI.class);
7172
// run the command immediately
7273
Mockito.doAnswer(invocation -> {
7374
Command command = invocation.getArgument(0);
@@ -88,6 +89,7 @@ public void setUp() throws IOException {
8889
public void handleDownloadRequest(DownloadEvent event) {
8990
}
9091
};
92+
handler.setTransferUI(ui);
9193
mockContext = Mockito.mock(TransferContext.class);
9294
Mockito.when(mockContext.contentLength()).thenReturn(TOTAL_BYTES);
9395
listener = Mockito.mock(TransferProgressListener.class);
@@ -193,6 +195,7 @@ public void handleDownloadRequest(DownloadEvent event) {
193195
successAtomic.set(success);
194196
});
195197

198+
customHandler.setTransferUI(ui);
196199
customHandler.handleDownloadRequest(downloadEvent);
197200

198201
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
@@ -47,6 +47,7 @@ public class ClassDownloadHandlerTest {
4747
private DownloadEvent downloadEvent;
4848
private OutputStream outputStream;
4949
private Element owner;
50+
private UI ui;
5051

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

58-
UI ui = Mockito.mock(UI.class);
59+
ui = Mockito.mock(UI.class);
5960
// run the command immediately
6061
Mockito.doAnswer(invocation -> {
6162
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: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ public class InputStreamDownloadHandlerTest {
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);
@@ -139,6 +140,7 @@ public void transferProgressListener_addListener_errorOccurred_errorListenerInvo
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(),
@@ -177,6 +179,7 @@ public void transferProgressListener_addListener_callbackIOExceptionOccurred_err
177179
Mockito.when(event.getSession()).thenReturn(session);
178180
Mockito.when(event.getResponse()).thenReturn(response);
179181
Mockito.when(event.getOwningElement()).thenReturn(owner);
182+
Mockito.when(event.getUI()).thenReturn(ui);
180183
OutputStream outputStreamMock = Mockito.mock(OutputStream.class);
181184
Mockito.when(event.getOutputStream()).thenReturn(outputStreamMock);
182185

@@ -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

@@ -242,6 +246,7 @@ public void transferProgressListener_addListener_callbackResponseError_errorList
242246
Mockito.when(event.getSession()).thenReturn(session);
243247
Mockito.when(event.getResponse()).thenReturn(response);
244248
Mockito.when(event.getOwningElement()).thenReturn(owner);
249+
Mockito.when(event.getUI()).thenReturn(ui);
245250
OutputStream outputStreamMock = Mockito.mock(OutputStream.class);
246251
Mockito.when(event.getOutputStream()).thenReturn(outputStreamMock);
247252

@@ -271,6 +276,7 @@ public void transferProgressListener_addListener_callbackResponseException_error
271276
Mockito.when(event.getSession()).thenReturn(session);
272277
Mockito.when(event.getResponse()).thenReturn(response);
273278
Mockito.when(event.getOwningElement()).thenReturn(owner);
279+
Mockito.when(event.getUI()).thenReturn(ui);
274280
OutputStream outputStreamMock = Mockito.mock(OutputStream.class);
275281
Mockito.when(event.getOutputStream()).thenReturn(outputStreamMock);
276282

@@ -322,6 +328,7 @@ public void inline_setFileNameInvokedByDefault() throws IOException {
322328
Mockito.when(event.getResponse()).thenReturn(response);
323329
Mockito.when(event.getOwningElement()).thenReturn(owner);
324330
Mockito.when(event.getOutputStream()).thenReturn(outputStream);
331+
Mockito.when(event.getUI()).thenReturn(ui);
325332
Mockito.when(response.getOutputStream()).thenReturn(outputStream);
326333
Mockito.when(response.getService()).thenReturn(service);
327334
Mockito.when(service.getMimeType(Mockito.anyString()))
@@ -348,6 +355,7 @@ public void attachment_doesNotSetFileNameWhenInlined() throws IOException {
348355
Mockito.when(event.getResponse()).thenReturn(response);
349356
Mockito.when(event.getOwningElement()).thenReturn(owner);
350357
Mockito.when(event.getOutputStream()).thenReturn(outputStream);
358+
Mockito.when(event.getUI()).thenReturn(ui);
351359
Mockito.when(response.getOutputStream()).thenReturn(outputStream);
352360
Mockito.when(response.getService()).thenReturn(service);
353361
Mockito.when(service.getMimeType(Mockito.anyString()))

0 commit comments

Comments
 (0)