From 461b6e0f7b3277fe980b7293a6ab644340af1d7c Mon Sep 17 00:00:00 2001
From: Aayush Atharva
Date: Thu, 8 Oct 2020 15:22:17 +0530
Subject: [PATCH] Add Close method for closing OutputStream and
PcapWriteHandler (#10638)
Motivation:
We should implement the Closeable method to properly close `OutputStream` and `PcapWriteHandler`. So whenever `handlerRemoved(ChannelHandlerContext)` is called or the user wants to stop the Pcap writes into `OutputStream`, we have a proper method to close it otherwise writing data to close `OutputStream` will result in `IOException`.
Modification:
Implemented `Closeable` in `PcapWriteHandler` which calls `PcapWriter#close` and closes `OutputStream` and stops Pcap writes.
Result:
Better handling of Pcap writes.
---
.../io/netty/handler/pcap/PcapHeaders.java | 2 -
.../netty/handler/pcap/PcapWriteHandler.java | 67 ++++++++++++++-----
.../io/netty/handler/pcap/PcapWriter.java | 33 +++++++--
3 files changed, 77 insertions(+), 25 deletions(-)
diff --git a/handler/src/main/java/io/netty/handler/pcap/PcapHeaders.java b/handler/src/main/java/io/netty/handler/pcap/PcapHeaders.java
index 96eacd27bf..a460222273 100644
--- a/handler/src/main/java/io/netty/handler/pcap/PcapHeaders.java
+++ b/handler/src/main/java/io/netty/handler/pcap/PcapHeaders.java
@@ -17,8 +17,6 @@ package io.netty.handler.pcap;
import io.netty.buffer.ByteBuf;
-import java.util.concurrent.TimeUnit;
-
final class PcapHeaders {
/**
diff --git a/handler/src/main/java/io/netty/handler/pcap/PcapWriteHandler.java b/handler/src/main/java/io/netty/handler/pcap/PcapWriteHandler.java
index 5c5e270651..f55dabd77d 100644
--- a/handler/src/main/java/io/netty/handler/pcap/PcapWriteHandler.java
+++ b/handler/src/main/java/io/netty/handler/pcap/PcapWriteHandler.java
@@ -31,6 +31,7 @@ import io.netty.util.internal.ObjectUtil;
import io.netty.util.internal.logging.InternalLogger;
import io.netty.util.internal.logging.InternalLoggerFactory;
+import java.io.Closeable;
import java.io.IOException;
import java.io.OutputStream;
import java.net.Inet4Address;
@@ -61,8 +62,11 @@ import java.net.InetSocketAddress;
*
*
*/
-public final class PcapWriteHandler extends ChannelDuplexHandler {
+public final class PcapWriteHandler extends ChannelDuplexHandler implements Closeable {
+ /**
+ * Logger for logging events
+ */
private final InternalLogger logger = InternalLoggerFactory.getInstance(PcapWriteHandler.class);
/**
@@ -108,12 +112,18 @@ public final class PcapWriteHandler extends ChannelDuplexHandler {
*/
private InetSocketAddress dstAddr;
+ /**
+ * Set to {@code true} if {@link #close()} is called and we should stop writing Pcap.
+ */
+ private boolean isClosed;
+
/**
* Create new {@link PcapWriteHandler} Instance.
* {@code captureZeroByte} is set to {@code false} and
* {@code writePcapGlobalHeader} is set to {@code true}.
*
- * @param outputStream OutputStream where Pcap data will be written
+ * @param outputStream OutputStream where Pcap data will be written. Call {@link #close()} to close this
+ * OutputStream.
* @throws NullPointerException If {@link OutputStream} is {@code null} then we'll throw an
* {@link NullPointerException}
*/
@@ -124,7 +134,8 @@ public final class PcapWriteHandler extends ChannelDuplexHandler {
/**
* Create new {@link PcapWriteHandler} Instance
*
- * @param outputStream OutputStream where Pcap data will be written
+ * @param outputStream OutputStream where Pcap data will be written. Call {@link #close()} to close this
+ * OutputStream.
* @param captureZeroByte Set to {@code true} to enable capturing packets with empty (0 bytes) payload.
* Otherwise, if set to {@code false}, empty packets will be filtered out.
* @param writePcapGlobalHeader Set to {@code true} to write Pcap Global Header on initialization.
@@ -214,24 +225,28 @@ public final class PcapWriteHandler extends ChannelDuplexHandler {
@Override
public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
- if (ctx.channel() instanceof SocketChannel) {
- handleTCP(ctx, msg, false);
- } else if (ctx.channel() instanceof DatagramChannel) {
- handleUDP(ctx, msg);
- } else {
- logger.debug("Discarding Pcap Write for Unknown Channel Type: {}", ctx.channel());
+ if (!isClosed) {
+ if (ctx.channel() instanceof SocketChannel) {
+ handleTCP(ctx, msg, false);
+ } else if (ctx.channel() instanceof DatagramChannel) {
+ handleUDP(ctx, msg);
+ } else {
+ logger.debug("Discarding Pcap Write for Unknown Channel Type: {}", ctx.channel());
+ }
}
super.channelRead(ctx, msg);
}
@Override
public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {
- if (ctx.channel() instanceof SocketChannel) {
- handleTCP(ctx, msg, true);
- } else if (ctx.channel() instanceof DatagramChannel) {
- handleUDP(ctx, msg);
- } else {
- logger.debug("Discarding Pcap Write for Unknown Channel Type: {}", ctx.channel());
+ if (!isClosed) {
+ if (ctx.channel() instanceof SocketChannel) {
+ handleTCP(ctx, msg, true);
+ } else if (ctx.channel() instanceof DatagramChannel) {
+ handleUDP(ctx, msg);
+ } else {
+ logger.debug("Discarding Pcap Write for Unknown Channel Type: {}", ctx.channel());
+ }
}
super.write(ctx, msg, promise);
}
@@ -495,7 +510,7 @@ public final class PcapWriteHandler extends ChannelDuplexHandler {
logger.debug("Finished Fake TCP FIN+ACK Flow to close connection");
}
- this.pCapWriter.close();
+ close();
super.handlerRemoved(ctx);
}
@@ -517,7 +532,25 @@ public final class PcapWriteHandler extends ChannelDuplexHandler {
logger.debug("Sent Fake TCP RST to close connection");
}
- this.pCapWriter.close();
+ close();
ctx.fireExceptionCaught(cause);
}
+
+ /**
+ * Close {@code PcapWriter} and {@link OutputStream}.
+ * Note: Calling this method does not close {@link PcapWriteHandler}.
+ * Only Pcap Writes are closed.
+ *
+ * @throws IOException If {@link OutputStream#close()} throws an exception
+ */
+ @Override
+ public void close() throws IOException {
+ if (isClosed) {
+ logger.debug("PcapWriterHandler is already closed");
+ } else {
+ isClosed = true;
+ pCapWriter.close();
+ logger.debug("PcapWriterHandler is now closed");
+ }
+ }
}
diff --git a/handler/src/main/java/io/netty/handler/pcap/PcapWriter.java b/handler/src/main/java/io/netty/handler/pcap/PcapWriter.java
index bae06b27b3..7933a8881a 100644
--- a/handler/src/main/java/io/netty/handler/pcap/PcapWriter.java
+++ b/handler/src/main/java/io/netty/handler/pcap/PcapWriter.java
@@ -16,19 +16,30 @@
package io.netty.handler.pcap;
import io.netty.buffer.ByteBuf;
+import io.netty.util.internal.logging.InternalLogger;
+import io.netty.util.internal.logging.InternalLoggerFactory;
import java.io.Closeable;
import java.io.IOException;
import java.io.OutputStream;
-import java.util.concurrent.TimeUnit;
final class PcapWriter implements Closeable {
+ /**
+ * Logger
+ */
+ private static final InternalLogger logger = InternalLoggerFactory.getInstance(PcapWriter.class);
+
/**
* {@link OutputStream} where we'll write Pcap data.
*/
private final OutputStream outputStream;
+ /**
+ * Set to {@code true} if {@link #outputStream} is closed.
+ */
+ private boolean isClosed;
+
/**
* This uses {@link OutputStream} for writing Pcap.
* Pcap Global Header is not written on construction.
@@ -58,12 +69,16 @@ final class PcapWriter implements Closeable {
* @throws IOException If {@link OutputStream#write(byte[])} throws an exception
*/
void writePacket(ByteBuf packetHeaderBuf, ByteBuf packet) throws IOException {
- long currentTime = System.currentTimeMillis();
+ if (isClosed) {
+ logger.debug("Pcap Write attempted on closed PcapWriter");
+ }
+
+ long timestamp = System.currentTimeMillis();
PcapHeaders.writePacketHeader(
packetHeaderBuf,
- (int) TimeUnit.MILLISECONDS.toSeconds(currentTime),
- (int) TimeUnit.MILLISECONDS.toMicros(currentTime) % 1000000,
+ (int) (timestamp / 1000L),
+ (int) (timestamp % 1000L * 1000L),
packet.readableBytes(),
packet.readableBytes()
);
@@ -74,7 +89,13 @@ final class PcapWriter implements Closeable {
@Override
public void close() throws IOException {
- outputStream.flush();
- outputStream.close();
+ if (isClosed) {
+ logger.debug("PcapWriter is already closed");
+ } else {
+ isClosed = true;
+ outputStream.flush();
+ outputStream.close();
+ logger.debug("PcapWriter is now closed");
+ }
}
}