1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19 package org.eclipse.jetty.websocket.common.extensions;
20
21 import static java.nio.file.StandardOpenOption.*;
22
23 import java.io.File;
24 import java.io.IOException;
25 import java.nio.ByteBuffer;
26 import java.nio.channels.SeekableByteChannel;
27 import java.nio.file.Files;
28 import java.nio.file.Path;
29 import java.util.Calendar;
30 import java.util.concurrent.atomic.AtomicInteger;
31
32 import org.eclipse.jetty.util.IO;
33 import org.eclipse.jetty.util.StringUtil;
34 import org.eclipse.jetty.util.log.Log;
35 import org.eclipse.jetty.util.log.Logger;
36 import org.eclipse.jetty.websocket.api.BatchMode;
37 import org.eclipse.jetty.websocket.api.WebSocketPolicy;
38 import org.eclipse.jetty.websocket.api.WriteCallback;
39 import org.eclipse.jetty.websocket.api.extensions.ExtensionConfig;
40 import org.eclipse.jetty.websocket.api.extensions.Frame;
41 import org.eclipse.jetty.websocket.common.Generator;
42 import org.eclipse.jetty.websocket.common.WebSocketFrame;
43
44 public class FrameCaptureExtension extends AbstractExtension
45 {
46 private static final Logger LOG = Log.getLogger(FrameCaptureExtension.class);
47
48 private static final int BUFSIZE = 32768;
49 private Generator generator;
50 private Path outputDir;
51 private String prefix = "frame";
52 private Path incomingFramesPath;
53 private Path outgoingFramesPath;
54
55 private AtomicInteger incomingCount = new AtomicInteger(0);
56 private AtomicInteger outgoingCount = new AtomicInteger(0);
57
58 private SeekableByteChannel incomingChannel;
59 private SeekableByteChannel outgoingChannel;
60
61 @Override
62 public String getName()
63 {
64 return "@frame-capture";
65 }
66
67 @Override
68 public void incomingFrame(Frame frame)
69 {
70 saveFrame(frame,false);
71 try
72 {
73 nextIncomingFrame(frame);
74 }
75 catch (Throwable t)
76 {
77 IO.close(incomingChannel);
78 incomingChannel = null;
79 throw t;
80 }
81 }
82
83 @Override
84 public void outgoingFrame(Frame frame, WriteCallback callback, BatchMode batchMode)
85 {
86 saveFrame(frame,true);
87 try
88 {
89 nextOutgoingFrame(frame,callback,batchMode);
90 }
91 catch (Throwable t)
92 {
93 IO.close(outgoingChannel);
94 outgoingChannel = null;
95 throw t;
96 }
97 }
98
99 private void saveFrame(Frame frame, boolean outgoing)
100 {
101 if (outputDir == null || generator == null)
102 {
103 return;
104 }
105
106 @SuppressWarnings("resource")
107 SeekableByteChannel channel = (outgoing) ? outgoingChannel : incomingChannel;
108
109 if (channel == null)
110 {
111 return;
112 }
113
114 ByteBuffer buf = getBufferPool().acquire(BUFSIZE,false);
115
116 try
117 {
118 WebSocketFrame f = WebSocketFrame.copy(frame);
119 f.setMasked(false);
120 generator.generateHeaderBytes(f,buf);
121 channel.write(buf);
122 if (frame.hasPayload())
123 {
124 channel.write(frame.getPayload().slice());
125 }
126 if (LOG.isDebugEnabled())
127 LOG.debug("Saved {} frame #{}",(outgoing) ? "outgoing" : "incoming",
128 (outgoing) ? outgoingCount.incrementAndGet() : incomingCount.incrementAndGet());
129 }
130 catch (IOException e)
131 {
132 LOG.warn("Unable to save frame: " + frame,e);
133 }
134 finally
135 {
136 getBufferPool().release(buf);
137 }
138 }
139
140 @Override
141 public void setConfig(ExtensionConfig config)
142 {
143 super.setConfig(config);
144
145 String cfgOutputDir = config.getParameter("output-dir",null);
146 if (StringUtil.isNotBlank(cfgOutputDir))
147 {
148 Path path = new File(cfgOutputDir).toPath();
149 if (Files.isDirectory(path) && Files.exists(path) && Files.isWritable(path))
150 {
151 this.outputDir = path;
152 }
153 else
154 {
155 LOG.warn("Unable to configure {}: not a valid output directory",path.toAbsolutePath().toString());
156 }
157 }
158
159 String cfgPrefix = config.getParameter("prefix","frame");
160 if (StringUtil.isNotBlank(cfgPrefix))
161 {
162 this.prefix = cfgPrefix;
163 }
164
165 if (this.outputDir != null)
166 {
167 try
168 {
169 Path dir = this.outputDir.toRealPath();
170
171
172 String tstamp = String.format("%1$tY%1$tm%1$td-%1$tH%1$tM%1$tS",Calendar.getInstance());
173 incomingFramesPath = dir.resolve(String.format("%s-%s-incoming.dat",this.prefix,tstamp));
174 outgoingFramesPath = dir.resolve(String.format("%s-%s-outgoing.dat",this.prefix,tstamp));
175
176 incomingChannel = Files.newByteChannel(incomingFramesPath,CREATE,WRITE);
177 outgoingChannel = Files.newByteChannel(outgoingFramesPath,CREATE,WRITE);
178
179 this.generator = new Generator(WebSocketPolicy.newServerPolicy(),getBufferPool(),false,true);
180 }
181 catch (IOException e)
182 {
183 LOG.warn("Unable to create capture file(s)",e);
184 }
185 }
186 }
187 }