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
74 public class MdcInjectionFilter extends CommonEventFilter {
75
76 public enum MdcKey {
77 handlerClass, remoteAddress, localAddress, remoteIp, remotePort, localIp, localPort
78 }
79
80
81 private static final AttributeKey CONTEXT_KEY = new AttributeKey(
82 MdcInjectionFilter.class, "context");
83
84 private ThreadLocal<Integer> callDepth = new ThreadLocal<Integer>() {
85 @Override
86 protected Integer initialValue() {
87 return 0;
88 }
89 };
90
91 private EnumSet<MdcKey> mdcKeys;
92
93
94
95
96
97
98
99
100 public MdcInjectionFilter(EnumSet<MdcKey> keys) {
101 this.mdcKeys = keys.clone();
102 }
103
104
105
106
107
108
109
110
111 public MdcInjectionFilter(MdcKey... keys) {
112 Set<MdcKey> keySet = new HashSet<MdcKey>(Arrays.asList(keys));
113 this.mdcKeys = EnumSet.copyOf(keySet);
114 }
115
116 public MdcInjectionFilter() {
117 this.mdcKeys = EnumSet.allOf(MdcKey.class);
118 }
119
120 @Override
121 protected void filter(IoFilterEvent event) throws Exception {
122
123
124 int currentCallDepth = callDepth.get();
125 callDepth.set(currentCallDepth + 1);
126 Map<String, String> context = getAndFillContext(event.getSession());
127
128 if (currentCallDepth == 0) {
129
130 for (Map.Entry<String, String> e : context.entrySet()) {
131 MDC.put(e.getKey(), e.getValue());
132 }
133 }
134
135 try {
136
137 event.fire();
138 } finally {
139 if (currentCallDepth == 0) {
140
141 for (String key : context.keySet()) {
142 MDC.remove(key);
143 }
144 callDepth.remove();
145 } else {
146 callDepth.set(currentCallDepth);
147 }
148 }
149 }
150
151 private Map<String, String> getAndFillContext(final IoSession session) {
152 Map<String, String> context = getContext(session);
153 if (context.isEmpty()) {
154 fillContext(session, context);
155 }
156 return context;
157 }
158
159 @SuppressWarnings("unchecked")
160 private static Map<String, String> getContext(final IoSession session) {
161 Map<String, String> context = (Map<String, String>) session.getAttribute(CONTEXT_KEY);
162 if (context == null) {
163 context = new ConcurrentHashMap<String, String>();
164 session.setAttribute(CONTEXT_KEY, context);
165 }
166 return context;
167 }
168
169
170
171
172
173
174
175 protected void fillContext(final IoSession session, final Map<String, String> context) {
176 if (mdcKeys.contains(MdcKey.handlerClass)) {
177 context.put(MdcKey.handlerClass.name(), session.getHandler()
178 .getClass().getName());
179 }
180 if (mdcKeys.contains(MdcKey.remoteAddress)) {
181 context.put(MdcKey.remoteAddress.name(), session.getRemoteAddress()
182 .toString());
183 }
184 if (mdcKeys.contains(MdcKey.localAddress)) {
185 context.put(MdcKey.localAddress.name(), session.getLocalAddress()
186 .toString());
187 }
188 if (session.getTransportMetadata().getAddressType() == InetSocketAddress.class) {
189 InetSocketAddress remoteAddress = (InetSocketAddress) session
190 .getRemoteAddress();
191 InetSocketAddress localAddress = (InetSocketAddress) session
192 .getLocalAddress();
193 if (mdcKeys.contains(MdcKey.remoteIp)) {
194 context.put(MdcKey.remoteIp.name(), remoteAddress.getAddress()
195 .getHostAddress());
196 }
197 if (mdcKeys.contains(MdcKey.remotePort)) {
198 context.put(MdcKey.remotePort.name(), String
199 .valueOf(remoteAddress.getPort()));
200 }
201 if (mdcKeys.contains(MdcKey.localIp)) {
202 context.put(MdcKey.localIp.name(), localAddress.getAddress()
203 .getHostAddress());
204 }
205 if (mdcKeys.contains(MdcKey.localPort)) {
206 context.put(MdcKey.localPort.name(), String
207 .valueOf(localAddress.getPort()));
208 }
209 }
210 }
211
212 public static String getProperty(IoSession session, String key) {
213 if (key == null) {
214 throw new NullPointerException("key should not be null");
215 }
216
217 Map<String, String> context = getContext(session);
218 String answer = context.get(key);
219 if (answer != null) {
220 return answer;
221 }
222
223 return MDC.get(key);
224 }
225
226
227
228
229
230
231
232
233
234 public static void setProperty(IoSession session, String key, String value) {
235 if (key == null) {
236 throw new NullPointerException("key should not be null");
237 }
238 if (value == null) {
239 removeProperty(session, key);
240 }
241 Map<String, String> context = getContext(session);
242 context.put(key, value);
243 MDC.put(key, value);
244 }
245
246 public static void removeProperty(IoSession session, String key) {
247 if (key == null) {
248 throw new NullPointerException("key should not be null");
249 }
250 Map<String, String> context = getContext(session);
251 context.remove(key);
252 MDC.remove(key);
253 }
254 }