Custom WebSocketAnnotationMethodMessageHandler

WebSocketAnnotationMethodMessageHandler.java

package org.springframework.web.socket.messaging;

import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.List;

import org.springframework.beans.factory.config.BeanExpressionContext;
import org.springframework.beans.factory.config.BeanExpressionResolver;
import org.springframework.beans.factory.config.ConfigurableBeanFactory;
import org.springframework.context.ConfigurableApplicationContext;
import org.springframework.core.annotation.AnnotationAwareOrderComparator;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.SubscribableChannel;
import org.springframework.messaging.handler.DestinationPatternsMessageCondition;
import org.springframework.messaging.handler.MessagingAdviceBean;
import org.springframework.messaging.handler.annotation.MessageMapping;
import org.springframework.messaging.handler.annotation.support.AnnotationExceptionHandlerMethodResolver;
import org.springframework.messaging.simp.SimpMessageMappingInfo;
import org.springframework.messaging.simp.SimpMessageSendingOperations;
import org.springframework.messaging.simp.SimpMessageTypeMessageCondition;
import org.springframework.messaging.simp.annotation.SubscribeMapping;
import org.springframework.messaging.simp.annotation.support.SimpAnnotationMethodMessageHandler;
import org.springframework.web.method.ControllerAdviceBean;

public class WebSocketAnnotationMethodMessageHandler extends SimpAnnotationMethodMessageHandler {

    public WebSocketAnnotationMethodMessageHandler(final SubscribableChannel clientInChannel, final MessageChannel clientOutChannel,
            final SimpMessageSendingOperations brokerTemplate) {
        super(clientInChannel, clientOutChannel, brokerTemplate);
    }

    @Override
    protected SimpMessageMappingInfo getMappingForMethod(final Method method, final Class<?> handlerType) {
        final MessageMapping typeAnnotation = AnnotationUtils.findAnnotation(handlerType, MessageMapping.class);
        final MessageMapping messageAnnotation = AnnotationUtils.findAnnotation(method, MessageMapping.class);
        if (messageAnnotation != null) {
            SimpMessageMappingInfo result = this.createMessageMappingCondition(messageAnnotation);
            if (typeAnnotation != null) {
                result = this.createMessageMappingCondition(typeAnnotation).combine(result);
            }
            return result;
        }
        final SubscribeMapping subscribeAnnotation = AnnotationUtils.findAnnotation(method, SubscribeMapping.class);
        if (subscribeAnnotation != null) {
            SimpMessageMappingInfo result = this.createSubscribeCondition(subscribeAnnotation);
            if (typeAnnotation != null) {
                result = this.createMessageMappingCondition(typeAnnotation).combine(result);
            }
            return result;
        }
        return null;
    }

    private SimpMessageMappingInfo createMessageMappingCondition(final MessageMapping annotation) {
        return new SimpMessageMappingInfo(SimpMessageTypeMessageCondition.MESSAGE, new DestinationPatternsMessageCondition(
                this.resolveDestinationNames(annotation.value()), this.getPathMatcher()));
    }

    private SimpMessageMappingInfo createSubscribeCondition(final SubscribeMapping annotation) {
        final SimpMessageTypeMessageCondition messageTypeMessageCondition = SimpMessageTypeMessageCondition.SUBSCRIBE;
        return new SimpMessageMappingInfo(messageTypeMessageCondition, new DestinationPatternsMessageCondition(
                this.resolveDestinationNames(annotation.value()), this.getPathMatcher()));
    }

    private String[] resolveDestinationNames(final String[] destinationNames) {
        final int length = destinationNames.length;
        final String[] result = new String[length];

        for (int i = 0; i < length; i++) {
            result[i] = this.resolveName(destinationNames[i]);
        }

        return result;
    }

    private String resolveName(final String name) {
        if (!(this.getApplicationContext() instanceof ConfigurableApplicationContext)) {
            return name;
        }

        final ConfigurableApplicationContext applicationContext = (ConfigurableApplicationContext) this.getApplicationContext();
        final ConfigurableBeanFactory configurableBeanFactory = applicationContext.getBeanFactory();

        final String placeholdersResolved = configurableBeanFactory.resolveEmbeddedValue(name);
        final BeanExpressionResolver exprResolver = configurableBeanFactory.getBeanExpressionResolver();
        if (exprResolver == null) {
            return name;
        }
        final Object result = exprResolver.evaluate(placeholdersResolved, new BeanExpressionContext(configurableBeanFactory, null));
        return result != null ? result.toString() : name;
    }

    @Override
    public void afterPropertiesSet() {
        this.initControllerAdviceCache();
        super.afterPropertiesSet();
    }

    private void initControllerAdviceCache() {
        if (this.getApplicationContext() == null) {
            return;
        }
        if (logger.isDebugEnabled()) {
            logger.debug("Looking for @MessageExceptionHandler mappings: " + this.getApplicationContext());
        }
        final List<ControllerAdviceBean> controllerAdvice = ControllerAdviceBean.findAnnotatedBeans(this.getApplicationContext());
        AnnotationAwareOrderComparator.sort(controllerAdvice);
        this.initMessagingAdviceCache(MessagingControllerAdviceBean.createFromList(controllerAdvice));
    }

    private void initMessagingAdviceCache(final List<MessagingAdviceBean> beans) {
        if (beans == null) {
            return;
        }
        for (final MessagingAdviceBean bean : beans) {
            final Class<?> beanType = bean.getBeanType();
            final AnnotationExceptionHandlerMethodResolver resolver = new AnnotationExceptionHandlerMethodResolver(beanType);
            if (resolver.hasExceptionMappings()) {
                this.registerExceptionHandlerAdvice(bean, resolver);
                logger.info("Detected @MessageExceptionHandler methods in " + bean);
            }
        }
    }

    /**
     * Adapt ControllerAdviceBean to MessagingAdviceBean.
     */
    private static class MessagingControllerAdviceBean implements MessagingAdviceBean {

        private final ControllerAdviceBean adviceBean;

        private MessagingControllerAdviceBean(final ControllerAdviceBean adviceBean) {
            this.adviceBean = adviceBean;
        }

        public static List<MessagingAdviceBean> createFromList(final List<ControllerAdviceBean> controllerAdvice) {
            final List<MessagingAdviceBean> messagingAdvice = new ArrayList<MessagingAdviceBean>(controllerAdvice.size());
            for (final ControllerAdviceBean bean : controllerAdvice) {
                messagingAdvice.add(new MessagingControllerAdviceBean(bean));
            }
            return messagingAdvice;
        }

        @Override
        public Class<?> getBeanType() {
            return adviceBean.getBeanType();
        }

        @Override
        public Object resolveBean() {
            return adviceBean.resolveBean();
        }

        @Override
        public boolean isApplicableToBeanType(final Class<?> beanType) {
            return adviceBean.isApplicableToBeanType(beanType);
        }

        @Override
        public int getOrder() {
            return adviceBean.getOrder();
        }
    }

}