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.messaging.support;
018
019import java.util.ArrayList;
020import java.util.Collections;
021import java.util.LinkedList;
022import java.util.List;
023import java.util.Map;
024
025import org.springframework.lang.Nullable;
026import org.springframework.messaging.Message;
027import org.springframework.util.Assert;
028import org.springframework.util.CollectionUtils;
029import org.springframework.util.LinkedMultiValueMap;
030import org.springframework.util.MultiValueMap;
031import org.springframework.util.ObjectUtils;
032
033/**
034 * {@link MessageHeaderAccessor} sub-class that supports storage and access of
035 * headers from an external source such as a message broker. Headers from the
036 * external source are kept separate from other headers, in a sub-map under the
037 * key {@link #NATIVE_HEADERS}. This allows separating processing headers from
038 * headers that need to be sent to or received from the external source.
039 *
040 * <p>This class is likely to be used through indirectly through a protocol
041 * specific sub-class that also provide factory methods to translate
042 * message headers to an from an external messaging source.
043 *
044 * @author Rossen Stoyanchev
045 * @since 4.0
046 */
047public class NativeMessageHeaderAccessor extends MessageHeaderAccessor {
048
049        /** The header name used to store native headers. */
050        public static final String NATIVE_HEADERS = "nativeHeaders";
051
052
053        /**
054         * Protected constructor to create a new instance.
055         */
056        protected NativeMessageHeaderAccessor() {
057                this((Map<String, List<String>>) null);
058        }
059
060        /**
061         * Protected constructor to create an instance with the given native headers.
062         * @param nativeHeaders native headers to create the message with (may be {@code null})
063         */
064        protected NativeMessageHeaderAccessor(@Nullable Map<String, List<String>> nativeHeaders) {
065                if (!CollectionUtils.isEmpty(nativeHeaders)) {
066                        setHeader(NATIVE_HEADERS, new LinkedMultiValueMap<>(nativeHeaders));
067                }
068        }
069
070        /**
071         * Protected constructor that copies headers from another message.
072         */
073        protected NativeMessageHeaderAccessor(@Nullable Message<?> message) {
074                super(message);
075                if (message != null) {
076                        @SuppressWarnings("unchecked")
077                        Map<String, List<String>> map = (Map<String, List<String>>) getHeader(NATIVE_HEADERS);
078                        if (map != null) {
079                                // setHeader checks for equality but we need copy of native headers
080                                setHeader(NATIVE_HEADERS, null);
081                                setHeader(NATIVE_HEADERS, new LinkedMultiValueMap<>(map));
082                        }
083                }
084        }
085
086
087        /**
088         * Subclasses can use this method to access the "native" headers sub-map.
089         */
090        @SuppressWarnings("unchecked")
091        @Nullable
092        protected Map<String, List<String>> getNativeHeaders() {
093                return (Map<String, List<String>>) getHeader(NATIVE_HEADERS);
094        }
095
096        /**
097         * Return a copy of the native headers sub-map, or an empty map.
098         */
099        public Map<String, List<String>> toNativeHeaderMap() {
100                Map<String, List<String>> map = getNativeHeaders();
101                return (map != null ? new LinkedMultiValueMap<>(map) : Collections.emptyMap());
102        }
103
104        @Override
105        public void setImmutable() {
106                if (isMutable()) {
107                        Map<String, List<String>> map = getNativeHeaders();
108                        if (map != null) {
109                                // setHeader checks for equality but we need immutable wrapper
110                                setHeader(NATIVE_HEADERS, null);
111                                setHeader(NATIVE_HEADERS, Collections.unmodifiableMap(map));
112                        }
113                        super.setImmutable();
114                }
115        }
116
117        @Override
118        public void copyHeaders(@Nullable Map<String, ?> headersToCopy) {
119                if (headersToCopy == null) {
120                        return;
121                }
122
123                @SuppressWarnings("unchecked")
124                Map<String, List<String>> map = (Map<String, List<String>>) headersToCopy.get(NATIVE_HEADERS);
125                if (map != null && map != getNativeHeaders()) {
126                        map.forEach(this::setNativeHeaderValues);
127                }
128
129                // setHeader checks for equality, native headers should be equal by now
130                super.copyHeaders(headersToCopy);
131        }
132
133        @Override
134        public void copyHeadersIfAbsent(@Nullable Map<String, ?> headersToCopy) {
135                if (headersToCopy == null) {
136                        return;
137                }
138
139                @SuppressWarnings("unchecked")
140                Map<String, List<String>> map = (Map<String, List<String>>) headersToCopy.get(NATIVE_HEADERS);
141                if (map != null && getNativeHeaders() == null) {
142                        map.forEach(this::setNativeHeaderValues);
143                }
144
145                super.copyHeadersIfAbsent(headersToCopy);
146        }
147
148        /**
149         * Whether the native header map contains the give header name.
150         * @param headerName the name of the header
151         */
152        public boolean containsNativeHeader(String headerName) {
153                Map<String, List<String>> map = getNativeHeaders();
154                return (map != null && map.containsKey(headerName));
155        }
156
157        /**
158         * Return all values for the specified native header, if present.
159         * @param headerName the name of the header
160         * @return the associated values, or {@code null} if none
161         */
162        @Nullable
163        public List<String> getNativeHeader(String headerName) {
164                Map<String, List<String>> map = getNativeHeaders();
165                return (map != null ? map.get(headerName) : null);
166        }
167
168        /**
169         * Return the first value for the specified native header, if present.
170         * @param headerName the name of the header
171         * @return the associated value, or {@code null} if none
172         */
173        @Nullable
174        public String getFirstNativeHeader(String headerName) {
175                Map<String, List<String>> map = getNativeHeaders();
176                if (map != null) {
177                        List<String> values = map.get(headerName);
178                        if (!CollectionUtils.isEmpty(values)) {
179                                return values.get(0);
180                        }
181                }
182                return null;
183        }
184
185        /**
186         * Set the specified native header value replacing existing values.
187         * <p>In order for this to work, the accessor must be {@link #isMutable()
188         * mutable}. See {@link MessageHeaderAccessor} for details.
189         */
190        public void setNativeHeader(String name, @Nullable String value) {
191                Assert.state(isMutable(), "Already immutable");
192                Map<String, List<String>> map = getNativeHeaders();
193                if (value == null) {
194                        if (map != null && map.get(name) != null) {
195                                setModified(true);
196                                map.remove(name);
197                        }
198                        return;
199                }
200                if (map == null) {
201                        map = new LinkedMultiValueMap<>(4);
202                        setHeader(NATIVE_HEADERS, map);
203                }
204                List<String> values = new LinkedList<>();
205                values.add(value);
206                if (!ObjectUtils.nullSafeEquals(values, getHeader(name))) {
207                        setModified(true);
208                        map.put(name, values);
209                }
210        }
211
212        /**
213         * Variant of {@link #addNativeHeader(String, String)} for all values.
214         * @since 5.2.12
215         */
216        public void setNativeHeaderValues(String name, @Nullable List<String> values) {
217                Assert.state(isMutable(), "Already immutable");
218                Map<String, List<String>> map = getNativeHeaders();
219                if (values == null) {
220                        if (map != null && map.get(name) != null) {
221                                setModified(true);
222                                map.remove(name);
223                        }
224                        return;
225                }
226                if (map == null) {
227                        map = new LinkedMultiValueMap<>(3);
228                        setHeader(NATIVE_HEADERS, map);
229                }
230                if (!ObjectUtils.nullSafeEquals(values, getHeader(name))) {
231                        setModified(true);
232                        map.put(name, new ArrayList<>(values));
233                }
234        }
235
236        /**
237         * Add the specified native header value to existing values.
238         * <p>In order for this to work, the accessor must be {@link #isMutable()
239         * mutable}. See {@link MessageHeaderAccessor} for details.
240         * @param name the name of the header
241         * @param value the header value to set
242         */
243        public void addNativeHeader(String name, @Nullable String value) {
244                Assert.state(isMutable(), "Already immutable");
245                if (value == null) {
246                        return;
247                }
248                Map<String, List<String>> nativeHeaders = getNativeHeaders();
249                if (nativeHeaders == null) {
250                        nativeHeaders = new LinkedMultiValueMap<>(4);
251                        setHeader(NATIVE_HEADERS, nativeHeaders);
252                }
253                List<String> values = nativeHeaders.computeIfAbsent(name, k -> new LinkedList<>());
254                values.add(value);
255                setModified(true);
256        }
257
258        /**
259         * Add the specified native headers to existing values.
260         * @param headers the headers to set
261         */
262        public void addNativeHeaders(@Nullable MultiValueMap<String, String> headers) {
263                if (headers == null) {
264                        return;
265                }
266                headers.forEach((key, values) -> values.forEach(value -> addNativeHeader(key, value)));
267        }
268
269        /**
270         * Remove the specified native header value replacing existing values.
271         * <p>In order for this to work, the accessor must be {@link #isMutable()
272         * mutable}. See {@link MessageHeaderAccessor} for details.
273         * @param headerName the name of the header
274         * @return the associated values, or {@code null} if the header was not present
275         */
276        @Nullable
277        public List<String> removeNativeHeader(String headerName) {
278                Assert.state(isMutable(), "Already immutable");
279                Map<String, List<String>> nativeHeaders = getNativeHeaders();
280                if (CollectionUtils.isEmpty(nativeHeaders)) {
281                        return null;
282                }
283                return nativeHeaders.remove(headerName);
284        }
285
286
287        /**
288         * Return the first value for the specified native header,
289         * or {@code null} if none.
290         * @param headerName the name of the header
291         * @param headers the headers map to introspect
292         * @return the associated value, or {@code null} if none
293         */
294        @SuppressWarnings("unchecked")
295        @Nullable
296        public static String getFirstNativeHeader(String headerName, Map<String, Object> headers) {
297                Map<String, List<String>> map = (Map<String, List<String>>) headers.get(NATIVE_HEADERS);
298                if (map != null) {
299                        List<String> values = map.get(headerName);
300                        if (!CollectionUtils.isEmpty(values)) {
301                                return values.get(0);
302                        }
303                }
304                return null;
305        }
306
307}