/*
    *  Licensed to the Apache Software Foundation (ASF) under one
    *  or more contributor license agreements.  See the NOTICE file
    *  distributed with this work for additional information
    *  regarding copyright ownership.  The ASF licenses this file
    *  to you under the Apache License, Version 2.0 (the
    *  "License"); you may not use this file except in compliance
    *  with the License.  You may obtain a copy of the License at
    *
   *    http://www.apache.org/licenses/LICENSE-2.0
   *
   *  Unless required by applicable law or agreed to in writing,
   *  software distributed under the License is distributed on an
   *  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
   *  KIND, either express or implied.  See the License for the
   *  specific language governing permissions and limitations
   *  under the License.
   *
   */
  package org.apache.mina.filter.logging;
  
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertNotNull;
  import static org.junit.Assert.assertNull;
  import static org.junit.Assert.fail;
  
  import java.io.IOException;
  import java.net.InetSocketAddress;
  import java.net.SocketAddress;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Set;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.TimeUnit;
  
  import org.apache.log4j.AppenderSkeleton;
  import org.apache.log4j.Level;
  import org.apache.log4j.spi.LoggingEvent;
  import org.apache.mina.core.buffer.IoBuffer;
  import org.apache.mina.core.filterchain.DefaultIoFilterChainBuilder;
  import org.apache.mina.core.filterchain.IoFilterAdapter;
  import org.apache.mina.core.future.ConnectFuture;
  import org.apache.mina.core.service.IoHandlerAdapter;
  import org.apache.mina.core.session.IdleStatus;
  import org.apache.mina.core.session.IoSession;
  import org.apache.mina.filter.codec.ProtocolCodecFactory;
  import org.apache.mina.filter.codec.ProtocolCodecFilter;
  import org.apache.mina.filter.codec.ProtocolDecoder;
  import org.apache.mina.filter.codec.ProtocolDecoderAdapter;
  import org.apache.mina.filter.codec.ProtocolDecoderOutput;
  import org.apache.mina.filter.codec.ProtocolEncoder;
  import org.apache.mina.filter.codec.ProtocolEncoderAdapter;
  import org.apache.mina.filter.codec.ProtocolEncoderOutput;
  import org.apache.mina.filter.executor.ExecutorFilter;
  import org.apache.mina.filter.statistic.ProfilerTimerFilter;
  import org.apache.mina.transport.socket.nio.NioSocketAcceptor;
  import org.apache.mina.transport.socket.nio.NioSocketConnector;
  import org.junit.After;
  import org.junit.Before;
  import org.junit.Test;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * Tests {@link MdcInjectionFilter} in various scenarios.
   *
   * @author <a href="http://mina.apache.org">Apache MINA Project</a>
   */
  public class MdcInjectionFilterTest {
  
      static Logger LOGGER = LoggerFactory.getLogger(MdcInjectionFilterTest.class);
  
      private static final int TIMEOUT = 5000;
  
      private final MyAppender appender = new MyAppender();
  
      private int port;
  
      private NioSocketAcceptor acceptor;
  
      private Level previousLevelRootLogger;
  
      private ExecutorFilter executorFilter1;
  
      private ExecutorFilter executorFilter2;
  
      @Before
      public void setUp() throws Exception {
          // comment out next line if you want to see normal logging
          org.apache.log4j.Logger.getRootLogger().removeAllAppenders();
          previousLevelRootLogger = org.apache.log4j.Logger.getRootLogger().getLevel();
          org.apache.log4j.Logger.getRootLogger().setLevel(Level.DEBUG);
          org.apache.log4j.Logger.getRootLogger().addAppender(appender);
          acceptor = new NioSocketAcceptor();
      }
  
     @After
     public void tearDown() throws Exception {
         acceptor.dispose(true);
         org.apache.log4j.Logger.getRootLogger().setLevel(previousLevelRootLogger);
 
         destroy(executorFilter1);
         destroy(executorFilter2);
 
         List<String> after = getThreadNames();
 
         // give acceptor some time to shut down
         Thread.sleep(50);
         after = getThreadNames();
 
         int count = 0;
 
         // NOTE: this is *not* intended to be a permanent fix for this test-case.
         // There used to be no API to block until the ExecutorService of AbstractIoService is terminated.
         // The API exists now : dispose(true) so we should get rid of this code.
 
         while (contains(after, "Nio") && count++ < 10) {
             Thread.sleep(50);
             after = getThreadNames();
         }
 
         while (contains(after, "pool") && count++ < 10) {
             Thread.sleep(50);
             after = getThreadNames();
         }
 
         // The problem is that we clear the events of the appender here, but it's possible that a thread from
         // a previous test still generates events during the execution of the next test
         appender.clear();
     }
 
     private void destroy(ExecutorFilter executorFilter) throws InterruptedException {
         if (executorFilter != null) {
             ExecutorService executor = (ExecutorService) executorFilter.getExecutor();
             executor.shutdown();
             while (!executor.isTerminated()) {
                 //System.out.println("Waiting for termination of " + executorFilter);  
                 executor.awaitTermination(10, TimeUnit.MILLISECONDS);
             }
         }
     }
 
     @Test
     public void testSimpleChain() throws IOException, InterruptedException {
         DefaultIoFilterChainBuilder chain = new DefaultIoFilterChainBuilder();
         chain.addFirst("mdc-injector", new MdcInjectionFilter());
         chain.addLast("dummy", new DummyIoFilter());
         chain.addLast("protocol", new ProtocolCodecFilter(new DummyProtocolCodecFactory()));
         test(chain);
     }
 
     @Test
     public void testExecutorFilterAtTheEnd() throws IOException, InterruptedException {
         executorFilter1 = new ExecutorFilter();
         DefaultIoFilterChainBuilder chain = new DefaultIoFilterChainBuilder();
         MdcInjectionFilter mdcInjectionFilter = new MdcInjectionFilter();
         chain.addFirst("mdc-injector1", mdcInjectionFilter);
         chain.addLast("dummy", new DummyIoFilter());
         chain.addLast("protocol", new ProtocolCodecFilter(new DummyProtocolCodecFactory()));
         chain.addLast("executor", executorFilter1);
         chain.addLast("mdc-injector2", mdcInjectionFilter);
         test(chain);
     }
 
     @Test
     public void testExecutorFilterAtBeginning() throws IOException, InterruptedException {
         executorFilter1 = new ExecutorFilter();
         DefaultIoFilterChainBuilder chain = new DefaultIoFilterChainBuilder();
         MdcInjectionFilter mdcInjectionFilter = new MdcInjectionFilter();
         chain.addLast("executor", executorFilter1);
         chain.addLast("mdc-injector", mdcInjectionFilter);
         chain.addLast("dummy", new DummyIoFilter());
         chain.addLast("protocol", new ProtocolCodecFilter(new DummyProtocolCodecFactory()));
         test(chain);
     }
 
     @Test
     public void testExecutorFilterBeforeProtocol() throws IOException, InterruptedException {
         executorFilter1 = new ExecutorFilter();
         DefaultIoFilterChainBuilder chain = new DefaultIoFilterChainBuilder();
         MdcInjectionFilter mdcInjectionFilter = new MdcInjectionFilter();
         chain.addLast("executor", executorFilter1);
         chain.addLast("mdc-injector", mdcInjectionFilter);
         chain.addLast("dummy", new DummyIoFilter());
         chain.addLast("protocol", new ProtocolCodecFilter(new DummyProtocolCodecFactory()));
         test(chain);
     }
 
     @Test
     public void testMultipleFilters() throws IOException, InterruptedException {
         executorFilter1 = new ExecutorFilter();
         DefaultIoFilterChainBuilder chain = new DefaultIoFilterChainBuilder();
         MdcInjectionFilter mdcInjectionFilter = new MdcInjectionFilter();
         chain.addLast("executor", executorFilter1);
         chain.addLast("mdc-injector", mdcInjectionFilter);
         chain.addLast("profiler", new ProfilerTimerFilter());
         chain.addLast("dummy", new DummyIoFilter());
         chain.addLast("logger", new LoggingFilter());
         chain.addLast("protocol", new ProtocolCodecFilter(new DummyProtocolCodecFactory()));
         test(chain);
     }
 
     @Test
     public void testTwoExecutorFilters() throws IOException, InterruptedException {
         DefaultIoFilterChainBuilder chain = new DefaultIoFilterChainBuilder();
         MdcInjectionFilter mdcInjectionFilter = new MdcInjectionFilter();
         executorFilter1 = new ExecutorFilter();
         executorFilter2 = new ExecutorFilter();
         chain.addLast("executorFilter1", executorFilter1);
         chain.addLast("mdc-injector1", mdcInjectionFilter);
         chain.addLast("protocol", new ProtocolCodecFilter(new DummyProtocolCodecFactory()));
         chain.addLast("dummy", new DummyIoFilter());
         chain.addLast("executorFilter2", executorFilter2);
         // add the MdcInjectionFilter instance after every ExecutorFilter
         // it's important to use the same MdcInjectionFilter instance
         chain.addLast("mdc-injector2", mdcInjectionFilter);
         test(chain);
     }
 
     @Test
     public void testOnlyRemoteAddress() throws IOException, InterruptedException {
         DefaultIoFilterChainBuilder chain = new DefaultIoFilterChainBuilder();
         chain.addFirst("mdc-injector", new MdcInjectionFilter(MdcInjectionFilter.MdcKey.remoteAddress));
         chain.addLast("dummy", new DummyIoFilter());
         chain.addLast("protocol", new ProtocolCodecFilter(new DummyProtocolCodecFactory()));
         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());
         connectAndWrite(connector, 0);
         connectAndWrite(connector, 1);
         // wait until Iohandler has received all events
         simpleIoHandler.messageSentLatch.await();
         simpleIoHandler.sessionIdleLatch.await();
         simpleIoHandler.sessionClosedLatch.await();
         connector.dispose(true);
 
         // 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 (event.getLoggerName().startsWith("org.apache.mina.core.service.AbstractIoService")) {
                 continue;
             }
             for (MdcInjectionFilter.MdcKey mdcKey : MdcInjectionFilter.MdcKey.values()) {
                 String key = mdcKey.name();
                 Object value = event.getMDC(key);
                 if (mdcKey == MdcInjectionFilter.MdcKey.remoteAddress) {
                     assertNotNull("MDC[remoteAddress] not set for [" + event.getMessage() + "]", value);
                 } else {
                     assertNull("MDC[" + key + "] set for [" + event.getMessage() + "]", value);
                 }
             }
         }
     }
 
     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(true);
 
         // make a copy to prevent ConcurrentModificationException
         List<LoggingEvent> events = new ArrayList<LoggingEvent>(appender.events);
 
         Set<String> loggersToCheck = new HashSet<String>();
         loggersToCheck.add(MdcInjectionFilterTest.class.getName());
         loggersToCheck.add(ProtocolCodecFilter.class.getName());
         loggersToCheck.add(LoggingFilter.class.getName());
 
         // verify that all logging events have correct MDC
         for (LoggingEvent event : events) {
 
             if (loggersToCheck.contains(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"));
             }
         }
         // assert 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);
 
         /**
          * Default constructor
          */
         public SimpleIoHandler() {
             super();
         }
 
         @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.closeNow();
         }
 
         @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 {
         /**
          * Default constructor
          */
         public DummyProtocolCodecFactory() {
             super();
         }
 
         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>());
 
         /**
          * Default constructor
          */
         public MyAppender() {
             super();
         }
 
         public void clear() {
             events.clear();
         }
 
         @Override
         protected void append(final LoggingEvent loggingEvent) {
             loggingEvent.getMDCCopy();
             events.add(loggingEvent);
         }
 
         public boolean requiresLayout() {
             return false;
         }
 
         public void close() {
             // Do nothing
         }
     }
 
     static class DummyIoFilter extends IoFilterAdapter {
         @Override
         public void sessionOpened(NextFilter nextFilter, IoSession session) throws Exception {
             LOGGER.info("DummyIoFilter.sessionOpened");
             nextFilter.sessionOpened(session);
         }
     }
 
     private List<String> getThreadNames() {
         List<String> list = new ArrayList<String>();
         int active = Thread.activeCount();
         Thread[] threads = new Thread[active];
         Thread.enumerate(threads);
         for (Thread thread : threads) {
             try {
                 String name = thread.getName();
                 list.add(name);
             } catch (NullPointerException ignore) {
             }
         }
         return list;
     }
 
     private boolean contains(List<String> list, String search) {
         for (String s : list) {
             if (s.contains(search)) {
                 return true;
             }
         }
         return false;
     }
 }