mdcinjectionfiltertest.java

mina是以Java实现的一个开源的网络程序框架

    }

    private void test(DefaultIoFilterChainBuilder chain) throws IOException, InterruptedException {
        // configure the server
        SimpleIoHandler simpleIoHandler = new SimpleIoHandler();
        acceptor.setHandler(simpleIoHandler);
        acceptor.bind(new InetSocketAddress(0));
        port = acceptor.getLocalAddress().getPort();
        acceptor.setFilterChainBuilder(chain);
        // create some clients
        NioSocketConnector connector = new NioSocketConnector();
        connector.setHandler(new IoHandlerAdapter());
        SocketAddress remoteAddressClients[] = new SocketAddress[2];
        remoteAddressClients[0] = connectAndWrite(connector,0);
        remoteAddressClients[1] = connectAndWrite(connector,1);
        // wait until Iohandler has received all events
        simpleIoHandler.messageSentLatch.await();
        simpleIoHandler.sessionIdleLatch.await();
        simpleIoHandler.sessionClosedLatch.await();
        connector.dispose();

        // make a copy to prevent ConcurrentModificationException
        List<LoggingEvent> events = new ArrayList<LoggingEvent>(appender.events);

        // verify that all logging events have correct MDC
        for (LoggingEvent event : events) {
            if (!ExecutorFilter.class.getName().equals(event.getLoggerName())) {
                Object remoteAddress = event.getMDC("remoteAddress");
                assertNotNull(
                    "MDC[remoteAddress] not set for [" + event.getMessage() + "]",
                    remoteAddress);
                assertNotNull(
                    "MDC[remotePort] not set for [" + event.getMessage() + "]",
                    event.getMDC("remotePort"));
                assertEquals(
                    "every event should have MDC[handlerClass]",
                    SimpleIoHandler.class.getName(),
                    event.getMDC("handlerClass") );
            }
        }
        // asert we have received all expected logging events for each client
        for (int i = 0; i < remoteAddressClients.length; i++) {
            SocketAddress remoteAddressClient = remoteAddressClients[i];
            assertEventExists(events, "sessionCreated", remoteAddressClient, null);
            assertEventExists(events, "sessionOpened", remoteAddressClient, null);
            assertEventExists(events, "decode", remoteAddressClient, null);
            assertEventExists(events, "messageReceived-1", remoteAddressClient, null);
            assertEventExists(events, "messageReceived-2", remoteAddressClient, "user-" + i);
            assertEventExists(events, "encode", remoteAddressClient, null);
            assertEventExists(events, "exceptionCaught", remoteAddressClient, "user-" + i);
            assertEventExists(events, "messageSent-1", remoteAddressClient, "user-" + i);
            assertEventExists(events, "messageSent-2", remoteAddressClient, null);
            assertEventExists(events, "sessionIdle", remoteAddressClient, "user-" + i);
            assertEventExists(events, "sessionClosed", remoteAddressClient, "user-" + i);
            assertEventExists(events, "sessionClosed", remoteAddressClient, "user-" + i);
            assertEventExists(events, "DummyIoFilter.sessionOpened", remoteAddressClient, "user-" + i);
        }
    }

    private SocketAddress connectAndWrite(NioSocketConnector connector, int clientNr) {
        ConnectFuture connectFuture = connector.connect(new InetSocketAddress("localhost", port));
        connectFuture.awaitUninterruptibly(TIMEOUT);
        IoBuffer message = IoBuffer.allocate(4).putInt(clientNr).flip();
        IoSession session = connectFuture.getSession();
        session.write(message).awaitUninterruptibly(TIMEOUT);
        return session.getLocalAddress();
    }

    private void assertEventExists(List<LoggingEvent> events,
                                   String message,
                                   SocketAddress address,
                                   String user) {
        InetSocketAddress remoteAddress = (InetSocketAddress) address;
        for (LoggingEvent event : events) {
            if (event.getMessage().equals(message) &&
                event.getMDC("remoteAddress").equals(remoteAddress.toString()) &&
                event.getMDC("remoteIp").equals(remoteAddress.getAddress().getHostAddress()) &&
                event.getMDC("remotePort").equals(remoteAddress.getPort()+"") ) {
                if (user == null && event.getMDC("user") == null) {
                    return;
                }
                if (user != null && user.equals(event.getMDC("user"))) {
                    return;
                }
                return;
            }
        }
        fail("No LoggingEvent found from [" + remoteAddress +"] with message [" + message + "]");
    }

    private static class SimpleIoHandler extends IoHandlerAdapter {

        CountDownLatch sessionIdleLatch = new CountDownLatch(2);
        CountDownLatch sessionClosedLatch = new CountDownLatch(2);
        CountDownLatch messageSentLatch = new CountDownLatch(2);

        @Override
        public void sessionCreated(IoSession session) throws Exception {
            logger.info("sessionCreated");
            session.getConfig().setIdleTime(IdleStatus.BOTH_IDLE, 1);
        }

        @Override
        public void sessionOpened(IoSession session) throws Exception {
            logger.info("sessionOpened");
        }

        @Override
        public void sessionClosed(IoSession session) throws Exception {
            logger.info("sessionClosed");
            sessionClosedLatch.countDown();
        }

        @Override
        public void sessionIdle(IoSession session, IdleStatus status) throws Exception {
            logger.info("sessionIdle");
            sessionIdleLatch.countDown();
            session.close(true);
        }

        @Override
        public void exceptionCaught(IoSession session, Throwable cause) throws Exception {
            logger.info("exceptionCaught", cause);
        }

        @Override
        public void messageReceived(IoSession session, Object message) throws Exception {
            logger.info("messageReceived-1");
            // adding a custom property to the context
            String user = "user-" + message;
            MdcInjectionFilter.setProperty(session, "user", user);
            logger.info("messageReceived-2");
            session.getService().broadcast(message);
            throw new RuntimeException("just a test, forcing exceptionCaught");
        }

        @Override
        public void messageSent(IoSession session, Object message) throws Exception {
            logger.info("messageSent-1");
            MdcInjectionFilter.removeProperty(session, "user");
            logger.info("messageSent-2");
            messageSentLatch.countDown();
        }
    }

    private static class DummyProtocolCodecFactory implements ProtocolCodecFactory {

        public ProtocolEncoder getEncoder(IoSession session) throws Exception {
            return new ProtocolEncoderAdapter() {
                public void encode(IoSession session, Object message, ProtocolEncoderOutput out) throws Exception {
                    logger.info("encode");
                    IoBuffer buffer = IoBuffer.allocate(4).putInt(123).flip();
                    out.write(buffer);
                }
            };
        }

        public ProtocolDecoder getDecoder(IoSession session) throws Exception {
            return new ProtocolDecoderAdapter() {
                public void decode(IoSession session, IoBuffer in, ProtocolDecoderOutput out) throws Exception {
                    if (in.remaining() >= 4) {
                        int value = in.getInt();
                        logger.info("decode");
                        out.write(value);
                    }
                }
            };
        }
    }

    private static class MyAppender extends AppenderSkeleton {

        List<LoggingEvent> events = Collections.synchronizedList(new ArrayList<LoggingEvent>());

        @Override
        protected void append(final LoggingEvent loggingEvent) {
            loggingEvent.getMDCCopy();
            events.add(loggingEvent);
        }

        @Override
        public boolean requiresLayout() {
            return false;
        }

        @Override
        public void close() {
        }
    }

    private static class DummyIoFilter extends IoFilterAdapter {
        @Override
        public void sessionOpened(NextFilter nextFilter, IoSession session) throws Exception {
            logger.info("DummyIoFilter.sessionOpened");
            nextFilter.sessionOpened(session);
        }
    }

}