001/*
002 * Copyright 2002-2019 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.socket.messaging;
018
019import java.util.ArrayList;
020import java.util.List;
021
022import org.springframework.context.ApplicationContext;
023import org.springframework.lang.Nullable;
024import org.springframework.messaging.MessageChannel;
025import org.springframework.messaging.SubscribableChannel;
026import org.springframework.messaging.handler.MessagingAdviceBean;
027import org.springframework.messaging.handler.annotation.support.AnnotationExceptionHandlerMethodResolver;
028import org.springframework.messaging.simp.SimpMessageSendingOperations;
029import org.springframework.messaging.simp.annotation.support.SimpAnnotationMethodMessageHandler;
030import org.springframework.web.method.ControllerAdviceBean;
031
032/**
033 * A sub-class of {@link SimpAnnotationMethodMessageHandler} to provide support
034 * for {@link org.springframework.web.bind.annotation.ControllerAdvice
035 * ControllerAdvice} with global {@code @MessageExceptionHandler} methods.
036 *
037 * @author Rossen Stoyanchev
038 * @since 4.2
039 */
040public class WebSocketAnnotationMethodMessageHandler extends SimpAnnotationMethodMessageHandler {
041
042        public WebSocketAnnotationMethodMessageHandler(SubscribableChannel clientInChannel,
043                        MessageChannel clientOutChannel, SimpMessageSendingOperations brokerTemplate) {
044
045                super(clientInChannel, clientOutChannel, brokerTemplate);
046        }
047
048
049        @Override
050        public void afterPropertiesSet() {
051                initControllerAdviceCache();
052                super.afterPropertiesSet();
053        }
054
055        private void initControllerAdviceCache() {
056                ApplicationContext context = getApplicationContext();
057                if (context == null) {
058                        return;
059                }
060                if (logger.isTraceEnabled()) {
061                        logger.trace("Looking for @MessageExceptionHandler mappings: " + context);
062                }
063                List<ControllerAdviceBean> beans = ControllerAdviceBean.findAnnotatedBeans(context);
064                initMessagingAdviceCache(MessagingControllerAdviceBean.createFromList(beans));
065        }
066
067        private void initMessagingAdviceCache(List<MessagingAdviceBean> beans) {
068                for (MessagingAdviceBean bean : beans) {
069                        Class<?> type = bean.getBeanType();
070                        if (type != null) {
071                                AnnotationExceptionHandlerMethodResolver resolver = new AnnotationExceptionHandlerMethodResolver(type);
072                                if (resolver.hasExceptionMappings()) {
073                                        registerExceptionHandlerAdvice(bean, resolver);
074                                        if (logger.isTraceEnabled()) {
075                                                logger.trace("Detected @MessageExceptionHandler methods in " + bean);
076                                        }
077                                }
078                        }
079                }
080        }
081
082
083        /**
084         * Adapt ControllerAdviceBean to MessagingAdviceBean.
085         */
086        private static final class MessagingControllerAdviceBean implements MessagingAdviceBean {
087
088                private final ControllerAdviceBean adviceBean;
089
090                private MessagingControllerAdviceBean(ControllerAdviceBean adviceBean) {
091                        this.adviceBean = adviceBean;
092                }
093
094                public static List<MessagingAdviceBean> createFromList(List<ControllerAdviceBean> beans) {
095                        List<MessagingAdviceBean> result = new ArrayList<>(beans.size());
096                        for (ControllerAdviceBean bean : beans) {
097                                result.add(new MessagingControllerAdviceBean(bean));
098                        }
099                        return result;
100                }
101
102                @Override
103                @Nullable
104                public Class<?> getBeanType() {
105                        return this.adviceBean.getBeanType();
106                }
107
108                @Override
109                public Object resolveBean() {
110                        return this.adviceBean.resolveBean();
111                }
112
113                @Override
114                public boolean isApplicableToBeanType(Class<?> beanType) {
115                        return this.adviceBean.isApplicableToBeanType(beanType);
116                }
117
118                @Override
119                public int getOrder() {
120                        return this.adviceBean.getOrder();
121                }
122        }
123
124}