001/*
002 *  Licensed to the Apache Software Foundation (ASF) under one
003 *  or more contributor license agreements.  See the NOTICE file
004 *  distributed with this work for additional information
005 *  regarding copyright ownership.  The ASF licenses this file
006 *  to you under the Apache License, Version 2.0 (the
007 *  "License"); you may not use this file except in compliance
008 *  with the License.  You may obtain a copy of the License at
009 *
010 *    http://www.apache.org/licenses/LICENSE-2.0
011 *
012 *  Unless required by applicable law or agreed to in writing,
013 *  software distributed under the License is distributed on an
014 *  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
015 *  KIND, either express or implied.  See the License for the
016 *  specific language governing permissions and limitations
017 *  under the License.
018 *
019 */
020package org.apache.mina.filter.logging;
021
022import java.net.InetSocketAddress;
023import java.util.EnumSet;
024import java.util.HashSet;
025import java.util.Map;
026import java.util.Set;
027import java.util.Arrays;
028import java.util.concurrent.ConcurrentHashMap;
029
030import org.apache.mina.core.filterchain.IoFilterEvent;
031import org.apache.mina.core.session.AttributeKey;
032import org.apache.mina.core.session.IoSession;
033import org.apache.mina.filter.util.CommonEventFilter;
034import org.slf4j.MDC;
035
036/**
037 * This filter will inject some key IoSession properties into the Mapped Diagnostic Context (MDC)
038 * <p>
039 * These properties will be set in the MDC for all logging events that are generated
040 * down the call stack, even in code that is not aware of MINA.
041 *
042 * By default, the following properties will be set for all transports:
043 * <ul>
044 *  <li>"handlerClass"</li>
045 *  <li>"remoteAddress"</li>
046 *  <li>"localAddress"</li>
047 * </ul>
048 *
049 * When <code>session.getTransportMetadata().getAddressType() == InetSocketAddress.class</code>
050 * the following properties will also be set:
051 * <ul>
052 * <li>"remoteIp"</li>
053 * <li>"remotePort"</li>
054 * <li>"localIp"</li>
055 * <li>"localPort"</li>
056 * </ul>
057 *
058 * User code can also add custom properties to the context, via {@link #setProperty(IoSession, String, String)}
059 *
060 * If you only want the MDC to be set for the IoHandler code, it's enough to add
061 * one MdcInjectionFilter at the end of the filter chain.
062 *
063 * If you want the MDC to be set for ALL code, you should
064 *   add an MdcInjectionFilter to the start of the chain
065 *   and add that same MdcInjectionFilter instance after EVERY ExecutorFilter in the chain
066 *
067 * Thus it's ok to have one instance of the MdcInjectionFilter and add it multiple times to the chain
068 * but you should avoid adding multiple instances to the chain.
069 *
070 * @author <a href="http://mina.apache.org">Apache MINA Project</a>
071 */
072
073public class MdcInjectionFilter extends CommonEventFilter {
074
075    public enum MdcKey {
076        handlerClass, remoteAddress, localAddress, remoteIp, remotePort, localIp, localPort
077    }
078
079    /** key used for storing the context map in the IoSession */
080    private static final AttributeKey CONTEXT_KEY = new AttributeKey(MdcInjectionFilter.class, "context");
081
082    private ThreadLocal<Integer> callDepth = new ThreadLocal<Integer>() {
083        @Override
084        protected Integer initialValue() {
085            return 0;
086        }
087    };
088
089    private EnumSet<MdcKey> mdcKeys;
090
091    /**
092     * Use this constructor when you want to specify which keys to add to the MDC.
093     * You could still add custom keys via {@link #setProperty(IoSession, String, String)}
094     * @param keys set of keys that should be added to the MDC
095     *
096     * @see #setProperty(org.apache.mina.core.session.IoSession, String, String)
097     */
098    public MdcInjectionFilter(EnumSet<MdcKey> keys) {
099        this.mdcKeys = keys.clone();
100    }
101
102    /**
103     * Use this constructor when you want to specify which keys to add to the MDC
104     * You could still add custom keys via {@link #setProperty(IoSession, String, String)}
105     * @param keys list of keys that should be added to the MDC
106     *
107     * @see #setProperty(org.apache.mina.core.session.IoSession, String, String)
108     */
109    public MdcInjectionFilter(MdcKey... keys) {
110        Set<MdcKey> keySet = new HashSet<MdcKey>(Arrays.asList(keys));
111        this.mdcKeys = EnumSet.copyOf(keySet);
112    }
113
114    public MdcInjectionFilter() {
115        this.mdcKeys = EnumSet.allOf(MdcKey.class);
116    }
117
118    @Override
119    protected void filter(IoFilterEvent event) throws Exception {
120        // since this method can potentially call into itself
121        // we need to check the call depth before clearing the MDC
122        int currentCallDepth = callDepth.get();
123        callDepth.set(currentCallDepth + 1);
124        Map<String, String> context = getAndFillContext(event.getSession());
125
126        if (currentCallDepth == 0) {
127            /* copy context to the MDC when necessary. */
128            for (Map.Entry<String, String> e : context.entrySet()) {
129                MDC.put(e.getKey(), e.getValue());
130            }
131        }
132
133        try {
134            /* propagate event down the filter chain */
135            event.fire();
136        } finally {
137            if (currentCallDepth == 0) {
138                /* remove context from the MDC */
139                for (String key : context.keySet()) {
140                    MDC.remove(key);
141                }
142                callDepth.remove();
143            } else {
144                callDepth.set(currentCallDepth);
145            }
146        }
147    }
148
149    private Map<String, String> getAndFillContext(final IoSession session) {
150        Map<String, String> context = getContext(session);
151        if (context.isEmpty()) {
152            fillContext(session, context);
153        }
154        return context;
155    }
156
157    @SuppressWarnings("unchecked")
158    private static Map<String, String> getContext(final IoSession session) {
159        Map<String, String> context = (Map<String, String>) session.getAttribute(CONTEXT_KEY);
160        if (context == null) {
161            context = new ConcurrentHashMap<String, String>();
162            session.setAttribute(CONTEXT_KEY, context);
163        }
164        return context;
165    }
166
167    /**
168     * write key properties of the session to the Mapped Diagnostic Context
169     * sub-classes could override this method to map more/other attributes
170     * @param session the session to map
171     * @param context key properties will be added to this map
172     */
173    protected void fillContext(final IoSession session, final Map<String, String> context) {
174        if (mdcKeys.contains(MdcKey.handlerClass)) {
175            context.put(MdcKey.handlerClass.name(), session.getHandler().getClass().getName());
176        }
177        if (mdcKeys.contains(MdcKey.remoteAddress)) {
178            context.put(MdcKey.remoteAddress.name(), session.getRemoteAddress().toString());
179        }
180        if (mdcKeys.contains(MdcKey.localAddress)) {
181            context.put(MdcKey.localAddress.name(), session.getLocalAddress().toString());
182        }
183        if (session.getTransportMetadata().getAddressType() == InetSocketAddress.class) {
184            InetSocketAddress remoteAddress = (InetSocketAddress) session.getRemoteAddress();
185            InetSocketAddress localAddress = (InetSocketAddress) session.getLocalAddress();
186
187            if (mdcKeys.contains(MdcKey.remoteIp)) {
188                context.put(MdcKey.remoteIp.name(), remoteAddress.getAddress().getHostAddress());
189            }
190            if (mdcKeys.contains(MdcKey.remotePort)) {
191                context.put(MdcKey.remotePort.name(), String.valueOf(remoteAddress.getPort()));
192            }
193            if (mdcKeys.contains(MdcKey.localIp)) {
194                context.put(MdcKey.localIp.name(), localAddress.getAddress().getHostAddress());
195            }
196            if (mdcKeys.contains(MdcKey.localPort)) {
197                context.put(MdcKey.localPort.name(), String.valueOf(localAddress.getPort()));
198            }
199        }
200    }
201
202    public static String getProperty(IoSession session, String key) {
203        if (key == null) {
204            throw new IllegalArgumentException("key should not be null");
205        }
206
207        Map<String, String> context = getContext(session);
208        String answer = context.get(key);
209        if (answer != null) {
210            return answer;
211        }
212
213        return MDC.get(key);
214    }
215
216    /**
217     * Add a property to the context for the given session
218     * This property will be added to the MDC for all subsequent events
219     * @param session The session for which you want to set a property
220     * @param key  The name of the property (should not be null)
221     * @param value The value of the property
222     */
223    public static void setProperty(IoSession session, String key, String value) {
224        if (key == null) {
225            throw new IllegalArgumentException("key should not be null");
226        }
227        if (value == null) {
228            removeProperty(session, key);
229        }
230        Map<String, String> context = getContext(session);
231        context.put(key, value);
232        MDC.put(key, value);
233    }
234
235    public static void removeProperty(IoSession session, String key) {
236        if (key == null) {
237            throw new IllegalArgumentException("key should not be null");
238        }
239        Map<String, String> context = getContext(session);
240        context.remove(key);
241        MDC.remove(key);
242    }
243}