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.Collections;
020import java.util.LinkedList;
021import java.util.List;
022import java.util.Map;
023
024import org.springframework.messaging.Message;
025import org.springframework.util.Assert;
026import org.springframework.util.CollectionUtils;
027import org.springframework.util.LinkedMultiValueMap;
028import org.springframework.util.MultiValueMap;
029import org.springframework.util.ObjectUtils;
030
031/**
032 * An extension of {@link MessageHeaderAccessor} that also stores and provides read/write
033 * access to message headers from an external source -- e.g. a Spring {@link Message}
034 * created to represent a STOMP message received from a STOMP client or message broker.
035 * Native message headers are kept in a {@code Map<String, List<String>>} under the key
036 * {@link #NATIVE_HEADERS}.
037 *
038 * <p>This class is not intended for direct use but is rather expected to be used
039 * indirectly through protocol-specific sub-classes such as
040 * {@link org.springframework.messaging.simp.stomp.StompHeaderAccessor StompHeaderAccessor}.
041 * Such sub-classes may provide factory methods to translate message headers from
042 * an external messaging source (e.g. STOMP) to Spring {@link Message} headers and
043 * reversely to translate Spring {@link Message} headers to a message to send to an
044 * external source.
045 *
046 * @author Rossen Stoyanchev
047 * @since 4.0
048 */
049public class NativeMessageHeaderAccessor extends MessageHeaderAccessor {
050
051        public static final String NATIVE_HEADERS = "nativeHeaders";
052
053
054        /**
055         * Protected constructor to create a new instance.
056         */
057        protected NativeMessageHeaderAccessor() {
058                this((Map<String, List<String>>) null);
059        }
060
061        /**
062         * Protected constructor to create an instance with the given native headers.
063         * @param nativeHeaders native headers to create the message with (may be {@code null})
064         */
065        protected NativeMessageHeaderAccessor(Map<String, List<String>> nativeHeaders) {
066                if (!CollectionUtils.isEmpty(nativeHeaders)) {
067                        setHeader(NATIVE_HEADERS, new LinkedMultiValueMap<String, String>(nativeHeaders));
068                }
069        }
070
071        /**
072         * Protected constructor that copies headers from another message.
073         */
074        protected NativeMessageHeaderAccessor(Message<?> message) {
075                super(message);
076                if (message != null) {
077                        @SuppressWarnings("unchecked")
078                        Map<String, List<String>> map = (Map<String, List<String>>) getHeader(NATIVE_HEADERS);
079                        if (map != null) {
080                                // Force removal since setHeader checks for equality
081                                removeHeader(NATIVE_HEADERS);
082                                setHeader(NATIVE_HEADERS, new LinkedMultiValueMap<String, String>(map));
083                        }
084                }
085        }
086
087
088        /**
089         * Subclasses can use this method to access the "native" headers sub-map.
090         */
091        @SuppressWarnings("unchecked")
092        private 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<String, String>(map) : Collections.<String, List<String>>emptyMap());
102        }
103
104        @Override
105        public void setImmutable() {
106                if (isMutable()) {
107                        Map<String, List<String>> map = getNativeHeaders();
108                        if (map != null) {
109                                // Force removal since setHeader checks for equality
110                                removeHeader(NATIVE_HEADERS);
111                                setHeader(NATIVE_HEADERS, Collections.<String, List<String>>unmodifiableMap(map));
112                        }
113                        super.setImmutable();
114                }
115        }
116
117        /**
118         * Whether the native header map contains the give header name.
119         * @param headerName the name of the header
120         */
121        public boolean containsNativeHeader(String headerName) {
122                Map<String, List<String>> map = getNativeHeaders();
123                return (map != null && map.containsKey(headerName));
124        }
125
126        /**
127         * Return all values for the specified native header, if present.
128         * @param headerName the name of the header
129         * @return the associated values, or {@code null} if none
130         */
131        public List<String> getNativeHeader(String headerName) {
132                Map<String, List<String>> map = getNativeHeaders();
133                return (map != null ? map.get(headerName) : null);
134        }
135
136        /**
137         * Return the first value for the specified native header, if present.
138         * @param headerName the name of the header
139         * @return the associated value, or {@code null} if none
140         */
141        public String getFirstNativeHeader(String headerName) {
142                Map<String, List<String>> map = getNativeHeaders();
143                if (map != null) {
144                        List<String> values = map.get(headerName);
145                        if (!CollectionUtils.isEmpty(values)) {
146                                return values.get(0);
147                        }
148                }
149                return null;
150        }
151
152        /**
153         * Set the specified native header value replacing existing values.
154         * <p>In order for this to work, the accessor must be {@link #isMutable()
155         * mutable}. See {@link MessageHeaderAccessor} for details.
156         */
157        public void setNativeHeader(String name, String value) {
158                Assert.state(isMutable(), "Already immutable");
159                Map<String, List<String>> map = getNativeHeaders();
160                if (value == null) {
161                        if (map != null && map.get(name) != null) {
162                                setModified(true);
163                                map.remove(name);
164                        }
165                        return;
166                }
167                if (map == null) {
168                        map = new LinkedMultiValueMap<String, String>(4);
169                        setHeader(NATIVE_HEADERS, map);
170                }
171                List<String> values = new LinkedList<String>();
172                values.add(value);
173                if (!ObjectUtils.nullSafeEquals(values, getHeader(name))) {
174                        setModified(true);
175                        map.put(name, values);
176                }
177        }
178
179        /**
180         * Add the specified native header value to existing values.
181         * <p>In order for this to work, the accessor must be {@link #isMutable()
182         * mutable}. See {@link MessageHeaderAccessor} for details.
183         * @param name the name of the header
184         * @param value the header value to set
185         */
186        public void addNativeHeader(String name, String value) {
187                Assert.state(isMutable(), "Already immutable");
188                if (value == null) {
189                        return;
190                }
191                Map<String, List<String>> nativeHeaders = getNativeHeaders();
192                if (nativeHeaders == null) {
193                        nativeHeaders = new LinkedMultiValueMap<String, String>(4);
194                        setHeader(NATIVE_HEADERS, nativeHeaders);
195                }
196                List<String> values = nativeHeaders.get(name);
197                if (values == null) {
198                        values = new LinkedList<String>();
199                        nativeHeaders.put(name, values);
200                }
201                values.add(value);
202                setModified(true);
203        }
204
205        /**
206         * Add the specified native headers to existing values.
207         * @param headers the headers to set
208         */
209        public void addNativeHeaders(MultiValueMap<String, String> headers) {
210                if (headers == null) {
211                        return;
212                }
213                for (Map.Entry<String, List<String>> headerEntry : headers.entrySet()) {
214                        for (String value : headerEntry.getValue()) {
215                                addNativeHeader(headerEntry.getKey(), value);
216                        }
217                }
218        }
219
220        /**
221         * Remove the specified native header value replacing existing values.
222         * <p>In order for this to work, the accessor must be {@link #isMutable()
223         * mutable}. See {@link MessageHeaderAccessor} for details.
224         * @param headerName the name of the header
225         * @return the associated values, or {@code null} if the header was not present
226         */
227        public List<String> removeNativeHeader(String headerName) {
228                Assert.state(isMutable(), "Already immutable");
229                Map<String, List<String>> nativeHeaders = getNativeHeaders();
230                if (CollectionUtils.isEmpty(nativeHeaders)) {
231                        return null;
232                }
233                return nativeHeaders.remove(headerName);
234        }
235
236
237        /**
238         * Return the first value for the specified native header,
239         * or {@code null} if none.
240         * @param headerName the name of the header
241         * @param headers the headers map to introspect
242         * @return the associated value, or {@code null} if none
243         */
244        @SuppressWarnings("unchecked")
245        public static String getFirstNativeHeader(String headerName, Map<String, Object> headers) {
246                Map<String, List<String>> map = (Map<String, List<String>>) headers.get(NATIVE_HEADERS);
247                if (map != null) {
248                        List<String> values = map.get(headerName);
249                        if (!CollectionUtils.isEmpty(values)) {
250                                return values.get(0);
251                        }
252                }
253                return null;
254        }
255
256}