001/*
002 * Copyright 2002-2019 the original author or authors.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *      https://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package org.springframework.web.servlet.mvc.method.annotation;
018
019import java.io.OutputStream;
020import java.util.concurrent.Callable;
021
022import javax.servlet.ServletRequest;
023import javax.servlet.http.HttpServletResponse;
024
025import org.springframework.core.MethodParameter;
026import org.springframework.core.ResolvableType;
027import org.springframework.http.ResponseEntity;
028import org.springframework.http.server.ServerHttpResponse;
029import org.springframework.http.server.ServletServerHttpResponse;
030import org.springframework.lang.Nullable;
031import org.springframework.util.Assert;
032import org.springframework.web.context.request.NativeWebRequest;
033import org.springframework.web.context.request.async.WebAsyncUtils;
034import org.springframework.web.filter.ShallowEtagHeaderFilter;
035import org.springframework.web.method.support.HandlerMethodReturnValueHandler;
036import org.springframework.web.method.support.ModelAndViewContainer;
037
038/**
039 * Supports return values of type
040 * {@link org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody}
041 * and also {@code ResponseEntity<StreamingResponseBody>}.
042 *
043 * @author Rossen Stoyanchev
044 * @since 4.2
045 */
046public class StreamingResponseBodyReturnValueHandler implements HandlerMethodReturnValueHandler {
047
048        @Override
049        public boolean supportsReturnType(MethodParameter returnType) {
050                if (StreamingResponseBody.class.isAssignableFrom(returnType.getParameterType())) {
051                        return true;
052                }
053                else if (ResponseEntity.class.isAssignableFrom(returnType.getParameterType())) {
054                        Class<?> bodyType = ResolvableType.forMethodParameter(returnType).getGeneric().resolve();
055                        return (bodyType != null && StreamingResponseBody.class.isAssignableFrom(bodyType));
056                }
057                return false;
058        }
059
060        @Override
061        @SuppressWarnings("resource")
062        public void handleReturnValue(@Nullable Object returnValue, MethodParameter returnType,
063                        ModelAndViewContainer mavContainer, NativeWebRequest webRequest) throws Exception {
064
065                if (returnValue == null) {
066                        mavContainer.setRequestHandled(true);
067                        return;
068                }
069
070                HttpServletResponse response = webRequest.getNativeResponse(HttpServletResponse.class);
071                Assert.state(response != null, "No HttpServletResponse");
072                ServerHttpResponse outputMessage = new ServletServerHttpResponse(response);
073
074                if (returnValue instanceof ResponseEntity) {
075                        ResponseEntity<?> responseEntity = (ResponseEntity<?>) returnValue;
076                        response.setStatus(responseEntity.getStatusCodeValue());
077                        outputMessage.getHeaders().putAll(responseEntity.getHeaders());
078                        returnValue = responseEntity.getBody();
079                        if (returnValue == null) {
080                                mavContainer.setRequestHandled(true);
081                                outputMessage.flush();
082                                return;
083                        }
084                }
085
086                ServletRequest request = webRequest.getNativeRequest(ServletRequest.class);
087                Assert.state(request != null, "No ServletRequest");
088                ShallowEtagHeaderFilter.disableContentCaching(request);
089
090                Assert.isInstanceOf(StreamingResponseBody.class, returnValue, "StreamingResponseBody expected");
091                StreamingResponseBody streamingBody = (StreamingResponseBody) returnValue;
092
093                Callable<Void> callable = new StreamingResponseBodyTask(outputMessage.getBody(), streamingBody);
094                WebAsyncUtils.getAsyncManager(webRequest).startCallableProcessing(callable, mavContainer);
095        }
096
097
098        private static class StreamingResponseBodyTask implements Callable<Void> {
099
100                private final OutputStream outputStream;
101
102                private final StreamingResponseBody streamingBody;
103
104                public StreamingResponseBodyTask(OutputStream outputStream, StreamingResponseBody streamingBody) {
105                        this.outputStream = outputStream;
106                        this.streamingBody = streamingBody;
107                }
108
109                @Override
110                public Void call() throws Exception {
111                        this.streamingBody.writeTo(this.outputStream);
112                        this.outputStream.flush();
113                        return null;
114                }
115        }
116
117}