001/*
002 * Copyright 2002-2017 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;
021import javax.servlet.ServletRequest;
022import javax.servlet.http.HttpServletResponse;
023
024import org.springframework.core.MethodParameter;
025import org.springframework.core.ResolvableType;
026import org.springframework.http.ResponseEntity;
027import org.springframework.http.server.ServerHttpResponse;
028import org.springframework.http.server.ServletServerHttpResponse;
029import org.springframework.util.Assert;
030import org.springframework.web.context.request.NativeWebRequest;
031import org.springframework.web.context.request.async.WebAsyncUtils;
032import org.springframework.web.filter.ShallowEtagHeaderFilter;
033import org.springframework.web.method.support.HandlerMethodReturnValueHandler;
034import org.springframework.web.method.support.ModelAndViewContainer;
035
036/**
037 * Supports return values of type
038 * {@link org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody}
039 * and also {@code ResponseEntity<StreamingResponseBody>}.
040 *
041 * @author Rossen Stoyanchev
042 * @since 4.2
043 */
044public class StreamingResponseBodyReturnValueHandler implements HandlerMethodReturnValueHandler {
045
046        @Override
047        public boolean supportsReturnType(MethodParameter returnType) {
048                if (StreamingResponseBody.class.isAssignableFrom(returnType.getParameterType())) {
049                        return true;
050                }
051                else if (ResponseEntity.class.isAssignableFrom(returnType.getParameterType())) {
052                        Class<?> bodyType = ResolvableType.forMethodParameter(returnType).getGeneric(0).resolve();
053                        return (bodyType != null && StreamingResponseBody.class.isAssignableFrom(bodyType));
054                }
055                return false;
056        }
057
058        @Override
059        public void handleReturnValue(Object returnValue, MethodParameter returnType,
060                        ModelAndViewContainer mavContainer, NativeWebRequest webRequest) throws Exception {
061
062                if (returnValue == null) {
063                        mavContainer.setRequestHandled(true);
064                        return;
065                }
066
067                HttpServletResponse response = webRequest.getNativeResponse(HttpServletResponse.class);
068                ServerHttpResponse outputMessage = new ServletServerHttpResponse(response);
069
070                if (returnValue instanceof ResponseEntity) {
071                        ResponseEntity<?> responseEntity = (ResponseEntity<?>) returnValue;
072                        response.setStatus(responseEntity.getStatusCodeValue());
073                        outputMessage.getHeaders().putAll(responseEntity.getHeaders());
074                        returnValue = responseEntity.getBody();
075                        if (returnValue == null) {
076                                mavContainer.setRequestHandled(true);
077                                outputMessage.flush();
078                                return;
079                        }
080                }
081
082                ServletRequest request = webRequest.getNativeRequest(ServletRequest.class);
083                ShallowEtagHeaderFilter.disableContentCaching(request);
084
085                Assert.isInstanceOf(StreamingResponseBody.class, returnValue, "StreamingResponseBody expected");
086                StreamingResponseBody streamingBody = (StreamingResponseBody) returnValue;
087
088                Callable<Void> callable = new StreamingResponseBodyTask(outputMessage.getBody(), streamingBody);
089                WebAsyncUtils.getAsyncManager(webRequest).startCallableProcessing(callable, mavContainer);
090        }
091
092
093        private static class StreamingResponseBodyTask implements Callable<Void> {
094
095                private final OutputStream outputStream;
096
097                private final StreamingResponseBody streamingBody;
098
099                public StreamingResponseBodyTask(OutputStream outputStream, StreamingResponseBody streamingBody) {
100                        this.outputStream = outputStream;
101                        this.streamingBody = streamingBody;
102                }
103
104                @Override
105                public Void call() throws Exception {
106                        this.streamingBody.writeTo(this.outputStream);
107                        return null;
108                }
109        }
110
111}