1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20 package org.apache.mina.filter.logging;
21
22 import java.net.InetSocketAddress;
23 import java.util.EnumSet;
24 import java.util.HashSet;
25 import java.util.Map;
26 import java.util.Set;
27 import java.util.Arrays;
28 import java.util.concurrent.ConcurrentHashMap;
29
30 import org.apache.mina.core.filterchain.IoFilterEvent;
31 import org.apache.mina.core.session.AttributeKey;
32 import org.apache.mina.core.session.IoSession;
33 import org.apache.mina.filter.util.CommonEventFilter;
34 import org.slf4j.MDC;
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73 public class MdcInjectionFilter extends CommonEventFilter {
74
75 public enum MdcKey {
76 handlerClass, remoteAddress, localAddress, remoteIp, remotePort, localIp, localPort
77 }
78
79
80 private static final AttributeKey CONTEXT_KEY = new AttributeKey(
81 MdcInjectionFilter.class, "context");
82
83 private ThreadLocal<Integer> callDepth = new ThreadLocal<Integer>() {
84 @Override
85 protected Integer initialValue() {
86 return 0;
87 }
88 };
89
90 private EnumSet<MdcKey> mdcKeys;
91
92
93
94
95
96
97
98
99 public MdcInjectionFilter(EnumSet<MdcKey> keys) {
100 this.mdcKeys = keys.clone();
101 }
102
103
104
105
106
107
108
109
110 public MdcInjectionFilter(MdcKey... keys) {
111 Set<MdcKey> keySet = new HashSet<MdcKey>(Arrays.asList(keys));
112 this.mdcKeys = EnumSet.copyOf(keySet);
113 }
114
115 public MdcInjectionFilter() {
116 this.mdcKeys = EnumSet.allOf(MdcKey.class);
117 }
118
119 @Override
120 protected void filter(IoFilterEvent event) throws Exception {
121
122
123 int currentCallDepth = callDepth.get();
124 callDepth.set(currentCallDepth + 1);
125 Map<String, String> context = getAndFillContext(event.getSession());
126
127 if (currentCallDepth == 0) {
128
129 for (Map.Entry<String, String> e : context.entrySet()) {
130 MDC.put(e.getKey(), e.getValue());
131 }
132 }
133
134 try {
135
136 event.fire();
137 } finally {
138 if (currentCallDepth == 0) {
139
140 for (String key : context.keySet()) {
141 MDC.remove(key);
142 }
143 callDepth.remove();
144 } else {
145 callDepth.set(currentCallDepth);
146 }
147 }
148 }
149
150 private Map<String, String> getAndFillContext(final IoSession session) {
151 Map<String, String> context = getContext(session);
152 if (context.isEmpty()) {
153 fillContext(session, context);
154 }
155 return context;
156 }
157
158 @SuppressWarnings("unchecked")
159 private static Map<String, String> getContext(final IoSession session) {
160 Map<String, String> context = (Map<String, String>) session.getAttribute(CONTEXT_KEY);
161 if (context == null) {
162 context = new ConcurrentHashMap<String, String>();
163 session.setAttribute(CONTEXT_KEY, context);
164 }
165 return context;
166 }
167
168
169
170
171
172
173
174 protected void fillContext(final IoSession session, final Map<String, String> context) {
175 if (mdcKeys.contains(MdcKey.handlerClass)) {
176 context.put(MdcKey.handlerClass.name(), session.getHandler()
177 .getClass().getName());
178 }
179 if (mdcKeys.contains(MdcKey.remoteAddress)) {
180 context.put(MdcKey.remoteAddress.name(), session.getRemoteAddress()
181 .toString());
182 }
183 if (mdcKeys.contains(MdcKey.localAddress)) {
184 context.put(MdcKey.localAddress.name(), session.getLocalAddress()
185 .toString());
186 }
187 if (session.getTransportMetadata().getAddressType() == InetSocketAddress.class) {
188 InetSocketAddress remoteAddress = (InetSocketAddress) session
189 .getRemoteAddress();
190 InetSocketAddress localAddress = (InetSocketAddress) session
191 .getLocalAddress();
192 if (mdcKeys.contains(MdcKey.remoteIp)) {
193 context.put(MdcKey.remoteIp.name(), remoteAddress.getAddress()
194 .getHostAddress());
195 }
196 if (mdcKeys.contains(MdcKey.remotePort)) {
197 context.put(MdcKey.remotePort.name(), String
198 .valueOf(remoteAddress.getPort()));
199 }
200 if (mdcKeys.contains(MdcKey.localIp)) {
201 context.put(MdcKey.localIp.name(), localAddress.getAddress()
202 .getHostAddress());
203 }
204 if (mdcKeys.contains(MdcKey.localPort)) {
205 context.put(MdcKey.localPort.name(), String
206 .valueOf(localAddress.getPort()));
207 }
208 }
209 }
210
211 public static String getProperty(IoSession session, String key) {
212 if (key == null) {
213 throw new NullPointerException("key should not be null");
214 }
215
216 Map<String, String> context = getContext(session);
217 String answer = context.get(key);
218 if (answer != null) {
219 return answer;
220 }
221
222 return MDC.get(key);
223 }
224
225
226
227
228
229
230
231
232
233 public static void setProperty(IoSession session, String key, String value) {
234 if (key == null) {
235 throw new NullPointerException("key should not be null");
236 }
237 if (value == null) {
238 removeProperty(session, key);
239 }
240 Map<String, String> context = getContext(session);
241 context.put(key, value);
242 MDC.put(key, value);
243 }
244
245 public static void removeProperty(IoSession session, String key) {
246 if (key == null) {
247 throw new NullPointerException("key should not be null");
248 }
249 Map<String, String> context = getContext(session);
250 context.remove(key);
251 MDC.remove(key);
252 }
253 }