only issue the FIN packet once the outgoing buffer has flushed

compscidr · Nov 7, 2024 · 43eb98d · 43eb98d
1 parent ecff78c
commit 43eb98d
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 0 deletions.
diff --git a/src/main/kotlin/com/jasonernst/kanonproxy/tcp/TcpSession.kt b/src/main/kotlin/com/jasonernst/kanonproxy/tcp/TcpSession.kt
@@ -33,6 +33,9 @@ abstract class TcpSession(
     private val logger = LoggerFactory.getLogger(javaClass)
     val isPsh = AtomicBoolean(false) // set when we have accepted a PSH packet
 
+    // set when teardown has been called but the outgoing buffer is not empty
+    val tearDownPending = AtomicBoolean(false)
+
     protected open val mtu =
         if (initialIpHeader == null) {
             logger.warn("Initial IP header is null, can't determine MTU")
@@ -65,6 +68,13 @@ abstract class TcpSession(
      */
     fun teardown(swapSourceAndDestination: Boolean = true): Packet? {
         logger.debug("Tcp session CLOSE function called in tcpState: ${tcpStateMachine.tcpState.value} swap?: $swapSourceAndDestination")
+
+        if (tcpStateMachine.outgoingBytesToSend() > 0) {
+            logger.debug("Outgoing bytes to send, setting teardown pending")
+            tearDownPending.set(true)
+            return null
+        }
+
         if (tcpStateMachine.transmissionControlBlock == null) {
             logger.debug("No TCB, returning to CLOSED")
             tcpStateMachine.tcpState.value = TcpState.CLOSED

diff --git a/src/main/kotlin/com/jasonernst/kanonproxy/tcp/TcpStateMachine.kt b/src/main/kotlin/com/jasonernst/kanonproxy/tcp/TcpStateMachine.kt
@@ -2610,6 +2610,14 @@ class TcpStateMachine(
         }
     }
 
+    fun outgoingBytesToSend(): Int {
+        return runBlocking {
+            tcbMutex.withLock {
+                return@runBlocking outgoingBuffer.position()
+            }
+        }
+    }
+
     /**
      * We assume that the buffer is positioned after the data to be encapsulated. Once it's been encapsulated, it's
      * removed from this buffer. The encapsulated packets are returned and also stored in the retransmit queue. Packets
@@ -2692,6 +2700,13 @@ class TcpStateMachine(
                         isPush = true
                         outgoingBuffer.clear()
                         logger.debug("ENC after clear position: ${outgoingBuffer.position()} limit: ${outgoingBuffer.limit()}")
+                        if (session.tearDownPending.get()) {
+                            logger.debug("TEARDOWN PENDING, sending FIN")
+                            val finPacket = session.teardown()
+                            if (finPacket != null) {
+                                packets.add(finPacket)
+                            }
+                        }
                     } else {
                         outgoingBuffer.compact()
                         logger.debug("ENC after compact position: ${outgoingBuffer.position()} limit: ${outgoingBuffer.limit()}")