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.json;
018
019import java.io.IOException;
020import java.lang.annotation.Annotation;
021import java.nio.charset.Charset;
022import java.util.ArrayList;
023import java.util.Collections;
024import java.util.HashMap;
025import java.util.List;
026import java.util.Map;
027
028import com.fasterxml.jackson.core.JsonEncoding;
029import com.fasterxml.jackson.core.JsonGenerator;
030import com.fasterxml.jackson.core.JsonProcessingException;
031import com.fasterxml.jackson.core.util.ByteArrayBuilder;
032import com.fasterxml.jackson.databind.JavaType;
033import com.fasterxml.jackson.databind.ObjectMapper;
034import com.fasterxml.jackson.databind.ObjectWriter;
035import com.fasterxml.jackson.databind.SequenceWriter;
036import com.fasterxml.jackson.databind.exc.InvalidDefinitionException;
037import org.reactivestreams.Publisher;
038import reactor.core.publisher.Flux;
039import reactor.core.publisher.Mono;
040
041import org.springframework.core.MethodParameter;
042import org.springframework.core.ResolvableType;
043import org.springframework.core.codec.CodecException;
044import org.springframework.core.codec.EncodingException;
045import org.springframework.core.codec.Hints;
046import org.springframework.core.io.buffer.DataBuffer;
047import org.springframework.core.io.buffer.DataBufferFactory;
048import org.springframework.core.log.LogFormatUtils;
049import org.springframework.http.MediaType;
050import org.springframework.http.codec.HttpMessageEncoder;
051import org.springframework.http.server.reactive.ServerHttpRequest;
052import org.springframework.http.server.reactive.ServerHttpResponse;
053import org.springframework.lang.Nullable;
054import org.springframework.util.Assert;
055import org.springframework.util.MimeType;
056
057/**
058 * Base class providing support methods for Jackson 2.9 encoding. For non-streaming use
059 * cases, {@link Flux} elements are collected into a {@link List} before serialization for
060 * performance reason.
061 *
062 * @author Sebastien Deleuze
063 * @author Arjen Poutsma
064 * @since 5.0
065 */
066public abstract class AbstractJackson2Encoder extends Jackson2CodecSupport implements HttpMessageEncoder<Object> {
067
068        private static final byte[] NEWLINE_SEPARATOR = {'\n'};
069
070        private static final Map<MediaType, byte[]> STREAM_SEPARATORS;
071
072        private static final Map<String, JsonEncoding> ENCODINGS;
073
074        static {
075                STREAM_SEPARATORS = new HashMap<>(4);
076                STREAM_SEPARATORS.put(MediaType.APPLICATION_STREAM_JSON, NEWLINE_SEPARATOR);
077                STREAM_SEPARATORS.put(MediaType.parseMediaType("application/stream+x-jackson-smile"), new byte[0]);
078
079                ENCODINGS = new HashMap<>(JsonEncoding.values().length + 1);
080                for (JsonEncoding encoding : JsonEncoding.values()) {
081                        ENCODINGS.put(encoding.getJavaName(), encoding);
082                }
083                ENCODINGS.put("US-ASCII", JsonEncoding.UTF8);
084        }
085
086
087        private final List<MediaType> streamingMediaTypes = new ArrayList<>(1);
088
089
090        /**
091         * Constructor with a Jackson {@link ObjectMapper} to use.
092         */
093        protected AbstractJackson2Encoder(ObjectMapper mapper, MimeType... mimeTypes) {
094                super(mapper, mimeTypes);
095        }
096
097
098        /**
099         * Configure "streaming" media types for which flushing should be performed
100         * automatically vs at the end of the stream.
101         * <p>By default this is set to {@link MediaType#APPLICATION_STREAM_JSON}.
102         * @param mediaTypes one or more media types to add to the list
103         * @see HttpMessageEncoder#getStreamingMediaTypes()
104         */
105        public void setStreamingMediaTypes(List<MediaType> mediaTypes) {
106                this.streamingMediaTypes.clear();
107                this.streamingMediaTypes.addAll(mediaTypes);
108        }
109
110
111        @Override
112        public boolean canEncode(ResolvableType elementType, @Nullable MimeType mimeType) {
113                Class<?> clazz = elementType.toClass();
114                if (!supportsMimeType(mimeType)) {
115                        return false;
116                }
117                if (mimeType != null && mimeType.getCharset() != null) {
118                        Charset charset = mimeType.getCharset();
119                        if (!ENCODINGS.containsKey(charset.name())) {
120                                return false;
121                        }
122                }
123                return (Object.class == clazz ||
124                                (!String.class.isAssignableFrom(elementType.resolve(clazz)) && getObjectMapper().canSerialize(clazz)));
125        }
126
127        @Override
128        public Flux<DataBuffer> encode(Publisher<?> inputStream, DataBufferFactory bufferFactory,
129                        ResolvableType elementType, @Nullable MimeType mimeType, @Nullable Map<String, Object> hints) {
130
131                Assert.notNull(inputStream, "'inputStream' must not be null");
132                Assert.notNull(bufferFactory, "'bufferFactory' must not be null");
133                Assert.notNull(elementType, "'elementType' must not be null");
134
135                if (inputStream instanceof Mono) {
136                        return Mono.from(inputStream)
137                                        .map(value -> encodeValue(value, bufferFactory, elementType, mimeType, hints))
138                                        .flux();
139                }
140                else {
141                        byte[] separator = streamSeparator(mimeType);
142                        if (separator != null) { // streaming
143                                try {
144                                        ObjectWriter writer = createObjectWriter(elementType, mimeType, hints);
145                                        ByteArrayBuilder byteBuilder = new ByteArrayBuilder(writer.getFactory()._getBufferRecycler());
146                                        JsonEncoding encoding = getJsonEncoding(mimeType);
147                                        JsonGenerator generator = getObjectMapper().getFactory().createGenerator(byteBuilder, encoding);
148                                        SequenceWriter sequenceWriter = writer.writeValues(generator);
149
150                                        return Flux.from(inputStream)
151                                                        .map(value -> encodeStreamingValue(value, bufferFactory, hints, sequenceWriter, byteBuilder,
152                                                                        separator))
153                                                        .doAfterTerminate(() -> {
154                                                                try {
155                                                                        byteBuilder.release();
156                                                                        generator.close();
157                                                                }
158                                                                catch (IOException ex) {
159                                                                        logger.error("Could not close Encoder resources", ex);
160                                                                }
161                                                        });
162                                }
163                                catch (IOException ex) {
164                                        return Flux.error(ex);
165                                }
166                        }
167                        else { // non-streaming
168                                ResolvableType listType = ResolvableType.forClassWithGenerics(List.class, elementType);
169                                return Flux.from(inputStream)
170                                                .collectList()
171                                                .map(list -> encodeValue(list, bufferFactory, listType, mimeType, hints))
172                                                .flux();
173                        }
174
175                }
176        }
177
178        @Override
179        public DataBuffer encodeValue(Object value, DataBufferFactory bufferFactory,
180                        ResolvableType valueType, @Nullable MimeType mimeType, @Nullable Map<String, Object> hints) {
181
182                ObjectWriter writer = createObjectWriter(valueType, mimeType, hints);
183                ByteArrayBuilder byteBuilder = new ByteArrayBuilder(writer.getFactory()._getBufferRecycler());
184                try {
185                        JsonEncoding encoding = getJsonEncoding(mimeType);
186
187                        logValue(hints, value);
188
189                        try (JsonGenerator generator = getObjectMapper().getFactory().createGenerator(byteBuilder, encoding)) {
190                                writer.writeValue(generator, value);
191                                generator.flush();
192                        }
193                        catch (InvalidDefinitionException ex) {
194                                throw new CodecException("Type definition error: " + ex.getType(), ex);
195                        }
196                        catch (JsonProcessingException ex) {
197                                throw new EncodingException("JSON encoding error: " + ex.getOriginalMessage(), ex);
198                        }
199                        catch (IOException ex) {
200                                throw new IllegalStateException("Unexpected I/O error while writing to byte array builder", ex);
201                        }
202
203                        byte[] bytes = byteBuilder.toByteArray();
204                        DataBuffer buffer = bufferFactory.allocateBuffer(bytes.length);
205                        buffer.write(bytes);
206
207                        return buffer;
208                }
209                finally {
210                        byteBuilder.release();
211                }
212        }
213
214        private DataBuffer encodeStreamingValue(Object value, DataBufferFactory bufferFactory, @Nullable Map<String, Object> hints,
215                        SequenceWriter sequenceWriter, ByteArrayBuilder byteArrayBuilder, byte[] separator) {
216
217                logValue(hints, value);
218
219                try {
220                        sequenceWriter.write(value);
221                        sequenceWriter.flush();
222                }
223                catch (InvalidDefinitionException ex) {
224                        throw new CodecException("Type definition error: " + ex.getType(), ex);
225                }
226                catch (JsonProcessingException ex) {
227                        throw new EncodingException("JSON encoding error: " + ex.getOriginalMessage(), ex);
228                }
229                catch (IOException ex) {
230                        throw new IllegalStateException("Unexpected I/O error while writing to byte array builder", ex);
231                }
232
233                byte[] bytes = byteArrayBuilder.toByteArray();
234                byteArrayBuilder.reset();
235
236                int offset;
237                int length;
238                if (bytes.length > 0 && bytes[0] == ' ') {
239                        // SequenceWriter writes an unnecessary space in between values
240                        offset = 1;
241                        length = bytes.length - 1;
242                }
243                else {
244                        offset = 0;
245                        length = bytes.length;
246                }
247                DataBuffer buffer = bufferFactory.allocateBuffer(length + separator.length);
248                buffer.write(bytes, offset, length);
249                buffer.write(separator);
250
251                return buffer;
252        }
253
254        private void logValue(@Nullable Map<String, Object> hints, Object value) {
255                if (!Hints.isLoggingSuppressed(hints)) {
256                        LogFormatUtils.traceDebug(logger, traceOn -> {
257                                String formatted = LogFormatUtils.formatValue(value, !traceOn);
258                                return Hints.getLogPrefix(hints) + "Encoding [" + formatted + "]";
259                        });
260                }
261        }
262
263        private ObjectWriter createObjectWriter(ResolvableType valueType, @Nullable MimeType mimeType,
264                        @Nullable Map<String, Object> hints) {
265
266                JavaType javaType = getJavaType(valueType.getType(), null);
267                Class<?> jsonView = (hints != null ? (Class<?>) hints.get(Jackson2CodecSupport.JSON_VIEW_HINT) : null);
268                ObjectWriter writer = (jsonView != null ?
269                                getObjectMapper().writerWithView(jsonView) : getObjectMapper().writer());
270
271                if (javaType.isContainerType()) {
272                        writer = writer.forType(javaType);
273                }
274
275                return customizeWriter(writer, mimeType, valueType, hints);
276        }
277
278        protected ObjectWriter customizeWriter(ObjectWriter writer, @Nullable MimeType mimeType,
279                        ResolvableType elementType, @Nullable Map<String, Object> hints) {
280
281                return writer;
282        }
283
284        @Nullable
285        private byte[] streamSeparator(@Nullable MimeType mimeType) {
286                for (MediaType streamingMediaType : this.streamingMediaTypes) {
287                        if (streamingMediaType.isCompatibleWith(mimeType)) {
288                                return STREAM_SEPARATORS.getOrDefault(streamingMediaType, NEWLINE_SEPARATOR);
289                        }
290                }
291                return null;
292        }
293
294        /**
295         * Determine the JSON encoding to use for the given mime type.
296         * @param mimeType the mime type as requested by the caller
297         * @return the JSON encoding to use (never {@code null})
298         * @since 5.0.5
299         */
300        protected JsonEncoding getJsonEncoding(@Nullable MimeType mimeType) {
301                if (mimeType != null && mimeType.getCharset() != null) {
302                        Charset charset = mimeType.getCharset();
303                        JsonEncoding result = ENCODINGS.get(charset.name());
304                        if (result != null) {
305                                return result;
306                        }
307                }
308                return JsonEncoding.UTF8;
309        }
310
311
312        // HttpMessageEncoder
313
314        @Override
315        public List<MimeType> getEncodableMimeTypes() {
316                return getMimeTypes();
317        }
318
319        @Override
320        public List<MediaType> getStreamingMediaTypes() {
321                return Collections.unmodifiableList(this.streamingMediaTypes);
322        }
323
324        @Override
325        public Map<String, Object> getEncodeHints(@Nullable ResolvableType actualType, ResolvableType elementType,
326                        @Nullable MediaType mediaType, ServerHttpRequest request, ServerHttpResponse response) {
327
328                return (actualType != null ? getHints(actualType) : Hints.none());
329        }
330
331
332        // Jackson2CodecSupport
333
334        @Override
335        protected <A extends Annotation> A getAnnotation(MethodParameter parameter, Class<A> annotType) {
336                return parameter.getMethodAnnotation(annotType);
337        }
338
339}