001/*
002 * Copyright 2002-2020 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.http.codec;
018
019import java.util.List;
020import java.util.Map;
021
022import org.apache.commons.logging.Log;
023import org.reactivestreams.Publisher;
024import reactor.core.publisher.Flux;
025import reactor.core.publisher.Mono;
026
027import org.springframework.core.ResolvableType;
028import org.springframework.core.codec.AbstractEncoder;
029import org.springframework.core.codec.Encoder;
030import org.springframework.core.codec.Hints;
031import org.springframework.core.io.buffer.DataBuffer;
032import org.springframework.core.io.buffer.DataBufferUtils;
033import org.springframework.core.io.buffer.PooledDataBuffer;
034import org.springframework.http.HttpLogging;
035import org.springframework.http.MediaType;
036import org.springframework.http.ReactiveHttpOutputMessage;
037import org.springframework.http.server.reactive.ServerHttpRequest;
038import org.springframework.http.server.reactive.ServerHttpResponse;
039import org.springframework.lang.Nullable;
040import org.springframework.util.Assert;
041import org.springframework.util.StringUtils;
042
043/**
044 * {@code HttpMessageWriter} that wraps and delegates to an {@link Encoder}.
045 *
046 * <p>Also a {@code HttpMessageWriter} that pre-resolves encoding hints
047 * from the extra information available on the server side such as the request
048 * or controller method annotations.
049 *
050 * @author Arjen Poutsma
051 * @author Sebastien Deleuze
052 * @author Rossen Stoyanchev
053 * @author Brian Clozel
054 * @author Sam Brannen
055 * @since 5.0
056 * @param <T> the type of objects in the input stream
057 */
058public class EncoderHttpMessageWriter<T> implements HttpMessageWriter<T> {
059
060        private final Encoder<T> encoder;
061
062        private final List<MediaType> mediaTypes;
063
064        @Nullable
065        private final MediaType defaultMediaType;
066
067
068        /**
069         * Create an instance wrapping the given {@link Encoder}.
070         */
071        public EncoderHttpMessageWriter(Encoder<T> encoder) {
072                Assert.notNull(encoder, "Encoder is required");
073                initLogger(encoder);
074                this.encoder = encoder;
075                this.mediaTypes = MediaType.asMediaTypes(encoder.getEncodableMimeTypes());
076                this.defaultMediaType = initDefaultMediaType(this.mediaTypes);
077        }
078
079        private static void initLogger(Encoder<?> encoder) {
080                if (encoder instanceof AbstractEncoder &&
081                                encoder.getClass().getName().startsWith("org.springframework.core.codec")) {
082                        Log logger = HttpLogging.forLog(((AbstractEncoder<?>) encoder).getLogger());
083                        ((AbstractEncoder<?>) encoder).setLogger(logger);
084                }
085        }
086
087        @Nullable
088        private static MediaType initDefaultMediaType(List<MediaType> mediaTypes) {
089                return mediaTypes.stream().filter(MediaType::isConcrete).findFirst().orElse(null);
090        }
091
092
093        /**
094         * Return the {@code Encoder} of this writer.
095         */
096        public Encoder<T> getEncoder() {
097                return this.encoder;
098        }
099
100        @Override
101        public List<MediaType> getWritableMediaTypes() {
102                return this.mediaTypes;
103        }
104
105
106        @Override
107        public boolean canWrite(ResolvableType elementType, @Nullable MediaType mediaType) {
108                return this.encoder.canEncode(elementType, mediaType);
109        }
110
111        @Override
112        public Mono<Void> write(Publisher<? extends T> inputStream, ResolvableType elementType,
113                        @Nullable MediaType mediaType, ReactiveHttpOutputMessage message, Map<String, Object> hints) {
114
115                MediaType contentType = updateContentType(message, mediaType);
116
117                Flux<DataBuffer> body = this.encoder.encode(
118                                inputStream, message.bufferFactory(), elementType, contentType, hints);
119
120                if (inputStream instanceof Mono) {
121                        return body
122                                        .singleOrEmpty()
123                                        .switchIfEmpty(Mono.defer(() -> {
124                                                message.getHeaders().setContentLength(0);
125                                                return message.setComplete().then(Mono.empty());
126                                        }))
127                                        .flatMap(buffer -> {
128                                                message.getHeaders().setContentLength(buffer.readableByteCount());
129                                                return message.writeWith(Mono.just(buffer)
130                                                                .doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release));
131                                        })
132                                        .doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release);
133                }
134
135                if (isStreamingMediaType(contentType)) {
136                        return message.writeAndFlushWith(body.map(buffer ->
137                                        Mono.just(buffer).doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release)));
138                }
139
140                return message.writeWith(body);
141        }
142
143        @Nullable
144        private MediaType updateContentType(ReactiveHttpOutputMessage message, @Nullable MediaType mediaType) {
145                MediaType result = message.getHeaders().getContentType();
146                if (result != null) {
147                        return result;
148                }
149                MediaType fallback = this.defaultMediaType;
150                result = (useFallback(mediaType, fallback) ? fallback : mediaType);
151                if (result != null) {
152                        result = addDefaultCharset(result, fallback);
153                        message.getHeaders().setContentType(result);
154                }
155                return result;
156        }
157
158        private static boolean useFallback(@Nullable MediaType main, @Nullable MediaType fallback) {
159                return (main == null || !main.isConcrete() ||
160                                main.equals(MediaType.APPLICATION_OCTET_STREAM) && fallback != null);
161        }
162
163        private static MediaType addDefaultCharset(MediaType main, @Nullable MediaType defaultType) {
164                if (main.getCharset() == null && defaultType != null && defaultType.getCharset() != null) {
165                        return new MediaType(main, defaultType.getCharset());
166                }
167                return main;
168        }
169
170        private boolean isStreamingMediaType(@Nullable MediaType mediaType) {
171                if (mediaType == null || !(this.encoder instanceof HttpMessageEncoder)) {
172                        return false;
173                }
174                for (MediaType streamingMediaType : ((HttpMessageEncoder<?>) this.encoder).getStreamingMediaTypes()) {
175                        if (mediaType.isCompatibleWith(streamingMediaType) && matchParameters(mediaType, streamingMediaType)) {
176                                return true;
177                        }
178                }
179                return false;
180        }
181
182        private boolean matchParameters(MediaType streamingMediaType, MediaType mediaType) {
183                for (String name : streamingMediaType.getParameters().keySet()) {
184                        String s1 = streamingMediaType.getParameter(name);
185                        String s2 = mediaType.getParameter(name);
186                        if (StringUtils.hasText(s1) && StringUtils.hasText(s2) && !s1.equalsIgnoreCase(s2)) {
187                                return false;
188                        }
189                }
190                return true;
191        }
192
193
194        // Server side only...
195
196        @Override
197        public Mono<Void> write(Publisher<? extends T> inputStream, ResolvableType actualType,
198                        ResolvableType elementType, @Nullable MediaType mediaType, ServerHttpRequest request,
199                        ServerHttpResponse response, Map<String, Object> hints) {
200
201                Map<String, Object> allHints = Hints.merge(hints,
202                                getWriteHints(actualType, elementType, mediaType, request, response));
203
204                return write(inputStream, elementType, mediaType, response, allHints);
205        }
206
207        /**
208         * Get additional hints for encoding for example based on the server request
209         * or annotations from controller method parameters. By default, delegate to
210         * the encoder if it is an instance of {@link HttpMessageEncoder}.
211         */
212        protected Map<String, Object> getWriteHints(ResolvableType streamType, ResolvableType elementType,
213                        @Nullable MediaType mediaType, ServerHttpRequest request, ServerHttpResponse response) {
214
215                if (this.encoder instanceof HttpMessageEncoder) {
216                        HttpMessageEncoder<?> encoder = (HttpMessageEncoder<?>) this.encoder;
217                        return encoder.getEncodeHints(streamType, elementType, mediaType, request, response);
218                }
219                return Hints.none();
220        }
221
222}