From fb4bd813c3c2c3b325a20636086d562886a4937d Mon Sep 17 00:00:00 2001 From: Jason Ernst Date: Thu, 19 Dec 2024 15:37:14 -0800 Subject: [PATCH 1/2] Added traffic accounting --- .../jasonernst/kanonproxy/DummyTrafficAccounting.kt | 11 +++++++++++ .../kotlin/com/jasonernst/kanonproxy/KAnonProxy.kt | 2 ++ .../main/kotlin/com/jasonernst/kanonproxy/Session.kt | 8 +++++++- .../com/jasonernst/kanonproxy/TrafficAccounting.kt | 6 ++++++ .../jasonernst/kanonproxy/tcp/AnonymousTcpSession.kt | 4 ++++ .../com/jasonernst/kanonproxy/tcp/TcpSession.kt | 3 +++ .../com/jasonernst/kanonproxy/udp/UdpSession.kt | 3 +++ .../kanonproxy/tcp/AnonymousTcpSessionTest.kt | 5 +++-- .../kotlin/com/jasonernst/kanonproxy/tcp/TcpClient.kt | 2 ++ .../kotlin/com/jasonernst/kanonproxy/ProxyServer.kt | 1 - 10 files changed, 41 insertions(+), 4 deletions(-) create mode 100644 core/src/main/kotlin/com/jasonernst/kanonproxy/DummyTrafficAccounting.kt create mode 100644 core/src/main/kotlin/com/jasonernst/kanonproxy/TrafficAccounting.kt diff --git a/core/src/main/kotlin/com/jasonernst/kanonproxy/DummyTrafficAccounting.kt b/core/src/main/kotlin/com/jasonernst/kanonproxy/DummyTrafficAccounting.kt new file mode 100644 index 0000000..36ba774 --- /dev/null +++ b/core/src/main/kotlin/com/jasonernst/kanonproxy/DummyTrafficAccounting.kt @@ -0,0 +1,11 @@ +package com.jasonernst.kanonproxy + +object DummyTrafficAccount : TrafficAccounting { + override fun recordToInternet(bytes: Long) { + // do nothing + } + + override fun recordFromInternet(bytes: Long) { + // do nothing + } +} \ No newline at end of file diff --git a/core/src/main/kotlin/com/jasonernst/kanonproxy/KAnonProxy.kt b/core/src/main/kotlin/com/jasonernst/kanonproxy/KAnonProxy.kt index b5f532d..03b7175 100644 --- a/core/src/main/kotlin/com/jasonernst/kanonproxy/KAnonProxy.kt +++ b/core/src/main/kotlin/com/jasonernst/kanonproxy/KAnonProxy.kt @@ -48,6 +48,7 @@ import java.util.concurrent.atomic.AtomicBoolean class KAnonProxy( val icmp: Icmp, val protector: VpnProtector = DummyProtector, + val trafficAccounting: TrafficAccounting = DummyTrafficAccount ) : SessionManager { private val logger = LoggerFactory.getLogger(javaClass) @@ -212,6 +213,7 @@ class KAnonProxy( protector, this, clientAddress, + trafficAccounting ) } if (isNewSession) { diff --git a/core/src/main/kotlin/com/jasonernst/kanonproxy/Session.kt b/core/src/main/kotlin/com/jasonernst/kanonproxy/Session.kt index 5d00091..a28e592 100644 --- a/core/src/main/kotlin/com/jasonernst/kanonproxy/Session.kt +++ b/core/src/main/kotlin/com/jasonernst/kanonproxy/Session.kt @@ -49,6 +49,7 @@ abstract class Session( val protector: VpnProtector, val sessionManager: SessionManager, val clientAddress: InetSocketAddress, + val trafficAccounting: TrafficAccounting = DummyTrafficAccount, ) { private val logger = LoggerFactory.getLogger(javaClass) abstract val channel: ByteChannel @@ -90,6 +91,7 @@ abstract class Session( protector: VpnProtector, sessionManager: SessionManager, clientAddress: InetSocketAddress, + trafficAccounting: TrafficAccounting, ): Session = when (initialIPHeader.protocol) { IpType.UDP.value -> { @@ -101,6 +103,7 @@ abstract class Session( protector, sessionManager, clientAddress, + trafficAccounting, ) } IpType.TCP.value -> { @@ -116,6 +119,7 @@ abstract class Session( protector, sessionManager, clientAddress, + trafficAccounting, ) } else -> { @@ -245,7 +249,8 @@ abstract class Session( val queue = outgoingQueue.take() logger.debug("Writing ${queue.limit()} bytes to remote channel") while (queue.hasRemaining()) { - channel.write(queue) + val bytesWritten = channel.write(queue) + trafficAccounting.recordToInternet(bytesWritten.toLong()) } } if (outgoingQueue.isNotEmpty()) { @@ -320,6 +325,7 @@ abstract class Session( val payload = ByteArray(len) readBuffer.get(payload, 0, len) logger.debug("Read {} bytes from {}", len, channel) + trafficAccounting.recordToInternet(len.toLong()) handlePayloadFromInternet(payload) readBuffer.clear() } diff --git a/core/src/main/kotlin/com/jasonernst/kanonproxy/TrafficAccounting.kt b/core/src/main/kotlin/com/jasonernst/kanonproxy/TrafficAccounting.kt new file mode 100644 index 0000000..dcb57f4 --- /dev/null +++ b/core/src/main/kotlin/com/jasonernst/kanonproxy/TrafficAccounting.kt @@ -0,0 +1,6 @@ +package com.jasonernst.kanonproxy + +interface TrafficAccounting { + fun recordToInternet(bytes: Long) + fun recordFromInternet(bytes: Long) +} \ No newline at end of file diff --git a/core/src/main/kotlin/com/jasonernst/kanonproxy/tcp/AnonymousTcpSession.kt b/core/src/main/kotlin/com/jasonernst/kanonproxy/tcp/AnonymousTcpSession.kt index b2e574f..dd15202 100644 --- a/core/src/main/kotlin/com/jasonernst/kanonproxy/tcp/AnonymousTcpSession.kt +++ b/core/src/main/kotlin/com/jasonernst/kanonproxy/tcp/AnonymousTcpSession.kt @@ -1,7 +1,9 @@ package com.jasonernst.kanonproxy.tcp import com.jasonernst.kanonproxy.ChangeRequest +import com.jasonernst.kanonproxy.DummyTrafficAccount import com.jasonernst.kanonproxy.SessionManager +import com.jasonernst.kanonproxy.TrafficAccounting import com.jasonernst.kanonproxy.VpnProtector import com.jasonernst.knet.Packet import com.jasonernst.knet.network.ip.IpHeader @@ -25,6 +27,7 @@ class AnonymousTcpSession( protector: VpnProtector, sessionManager: SessionManager, clientAddress: InetSocketAddress, + trafficAccounting: TrafficAccounting ) : TcpSession( initialIpHeader = initialIpHeader, initialTransportHeader = initialTransportHeader, @@ -33,6 +36,7 @@ class AnonymousTcpSession( protector = protector, sessionManager = sessionManager, clientAddress = clientAddress, + trafficAccounting = trafficAccounting, ) { companion object { const val CONNECTION_POLL_MS: Long = 500 diff --git a/core/src/main/kotlin/com/jasonernst/kanonproxy/tcp/TcpSession.kt b/core/src/main/kotlin/com/jasonernst/kanonproxy/tcp/TcpSession.kt index 9c90b7f..62e9ee0 100644 --- a/core/src/main/kotlin/com/jasonernst/kanonproxy/tcp/TcpSession.kt +++ b/core/src/main/kotlin/com/jasonernst/kanonproxy/tcp/TcpSession.kt @@ -2,6 +2,7 @@ package com.jasonernst.kanonproxy.tcp import com.jasonernst.kanonproxy.Session import com.jasonernst.kanonproxy.SessionManager +import com.jasonernst.kanonproxy.TrafficAccounting import com.jasonernst.kanonproxy.VpnProtector import com.jasonernst.knet.Packet import com.jasonernst.knet.network.ip.IpHeader @@ -24,6 +25,7 @@ abstract class TcpSession( protector: VpnProtector, sessionManager: SessionManager, clientAddress: InetSocketAddress, + trafficAccounting: TrafficAccounting, ) : Session( initialIpHeader = initialIpHeader, initialTransportHeader = initialTransportHeader, @@ -32,6 +34,7 @@ abstract class TcpSession( protector = protector, sessionManager = sessionManager, clientAddress = clientAddress, + trafficAccounting = trafficAccounting, ) { private val logger = LoggerFactory.getLogger(javaClass) val isPsh = AtomicBoolean(false) // set when we have accepted a PSH packet diff --git a/core/src/main/kotlin/com/jasonernst/kanonproxy/udp/UdpSession.kt b/core/src/main/kotlin/com/jasonernst/kanonproxy/udp/UdpSession.kt index a95b1a6..b6bdf39 100644 --- a/core/src/main/kotlin/com/jasonernst/kanonproxy/udp/UdpSession.kt +++ b/core/src/main/kotlin/com/jasonernst/kanonproxy/udp/UdpSession.kt @@ -3,6 +3,7 @@ package com.jasonernst.kanonproxy.udp import com.jasonernst.kanonproxy.ChangeRequest import com.jasonernst.kanonproxy.Session import com.jasonernst.kanonproxy.SessionManager +import com.jasonernst.kanonproxy.TrafficAccounting import com.jasonernst.kanonproxy.VpnProtector import com.jasonernst.knet.Packet import com.jasonernst.knet.network.ip.IpHeader @@ -30,6 +31,7 @@ class UdpSession( protector: VpnProtector, sessionManager: SessionManager, clientAddress: InetSocketAddress, + trafficAccounting: TrafficAccounting ) : Session( initialIpHeader = initialIpHeader, initialTransportHeader = initialTransportHeader, @@ -38,6 +40,7 @@ class UdpSession( protector = protector, sessionManager = sessionManager, clientAddress = clientAddress, + trafficAccounting = trafficAccounting ) { private val logger = LoggerFactory.getLogger(javaClass) diff --git a/core/src/test/kotlin/com/jasonernst/kanonproxy/tcp/AnonymousTcpSessionTest.kt b/core/src/test/kotlin/com/jasonernst/kanonproxy/tcp/AnonymousTcpSessionTest.kt index eaed8da..ae9de27 100644 --- a/core/src/test/kotlin/com/jasonernst/kanonproxy/tcp/AnonymousTcpSessionTest.kt +++ b/core/src/test/kotlin/com/jasonernst/kanonproxy/tcp/AnonymousTcpSessionTest.kt @@ -1,6 +1,7 @@ package com.jasonernst.kanonproxy.tcp import com.jasonernst.kanonproxy.DummyProtector +import com.jasonernst.kanonproxy.DummyTrafficAccount import com.jasonernst.kanonproxy.SessionManager import com.jasonernst.knet.Packet import com.jasonernst.knet.network.ip.IpType @@ -32,7 +33,7 @@ class AnonymousTcpSessionTest { val sessionManager = mockk() every { sessionManager.isRunning() } returns true - val session = AnonymousTcpSession(ipHeader, tcpHeader, ByteArray(0), returnQueue, DummyProtector, sessionManager, clientAddress) + val session = AnonymousTcpSession(ipHeader, tcpHeader, ByteArray(0), returnQueue, DummyProtector, sessionManager, clientAddress, DummyTrafficAccount) // wait until its connecting while (session.isConnecting.get().not()) { @@ -53,7 +54,7 @@ class AnonymousTcpSessionTest { ) val tcpHeader2 = TcpHeader(syn = true, destinationPort = 80u) val returnQueue2 = LinkedBlockingDeque() - val session2 = AnonymousTcpSession(ipHeader2, tcpHeader2, ByteArray(0), returnQueue2, DummyProtector, sessionManager, clientAddress) + val session2 = AnonymousTcpSession(ipHeader2, tcpHeader2, ByteArray(0), returnQueue2, DummyProtector, sessionManager, clientAddress, DummyTrafficAccount) // wait until its connecting while (session2.isConnecting.get().not()) { diff --git a/core/src/test/kotlin/com/jasonernst/kanonproxy/tcp/TcpClient.kt b/core/src/test/kotlin/com/jasonernst/kanonproxy/tcp/TcpClient.kt index 1768848..611945d 100644 --- a/core/src/test/kotlin/com/jasonernst/kanonproxy/tcp/TcpClient.kt +++ b/core/src/test/kotlin/com/jasonernst/kanonproxy/tcp/TcpClient.kt @@ -3,6 +3,7 @@ package com.jasonernst.kanonproxy.tcp import com.jasonernst.icmp.common.v4.IcmpV4DestinationUnreachablePacket import com.jasonernst.icmp.common.v6.IcmpV6DestinationUnreachablePacket import com.jasonernst.kanonproxy.BidirectionalByteChannel +import com.jasonernst.kanonproxy.DummyTrafficAccount import com.jasonernst.kanonproxy.KAnonProxy import com.jasonernst.knet.Packet import com.jasonernst.knet.SentinelPacket @@ -54,6 +55,7 @@ class TcpClient( mockk(relaxed = true), mockk(relaxed = true), clientAddress, + DummyTrafficAccount ) { private val clientId = UUID.randomUUID() private val logger = LoggerFactory.getLogger(javaClass) diff --git a/server/src/main/kotlin/com/jasonernst/kanonproxy/ProxyServer.kt b/server/src/main/kotlin/com/jasonernst/kanonproxy/ProxyServer.kt index 9ef46e1..b211364 100644 --- a/server/src/main/kotlin/com/jasonernst/kanonproxy/ProxyServer.kt +++ b/server/src/main/kotlin/com/jasonernst/kanonproxy/ProxyServer.kt @@ -139,7 +139,6 @@ class ProxyServer( val keyStream = selectedKeys.parallelStream() keyStream.forEach { if (it.isReadable && it.isValid) { - logger.debug("READ") readFromClient() } if (it.isWritable && it.isValid) { From 7d77fb3f4b91e6a3710c61342dbfb001c02c7357 Mon Sep 17 00:00:00 2001 From: Jason Ernst Date: Thu, 19 Dec 2024 15:37:32 -0800 Subject: [PATCH 2/2] lint fix --- .../kanonproxy/DummyTrafficAccounting.kt | 2 +- .../com/jasonernst/kanonproxy/KAnonProxy.kt | 4 ++-- .../kanonproxy/TrafficAccounting.kt | 3 ++- .../kanonproxy/tcp/AnonymousTcpSession.kt | 3 +-- .../jasonernst/kanonproxy/udp/UdpSession.kt | 4 ++-- .../kanonproxy/tcp/AnonymousTcpSessionTest.kt | 24 +++++++++++++++++-- .../jasonernst/kanonproxy/tcp/TcpClient.kt | 2 +- 7 files changed, 31 insertions(+), 11 deletions(-) diff --git a/core/src/main/kotlin/com/jasonernst/kanonproxy/DummyTrafficAccounting.kt b/core/src/main/kotlin/com/jasonernst/kanonproxy/DummyTrafficAccounting.kt index 36ba774..1baa1ba 100644 --- a/core/src/main/kotlin/com/jasonernst/kanonproxy/DummyTrafficAccounting.kt +++ b/core/src/main/kotlin/com/jasonernst/kanonproxy/DummyTrafficAccounting.kt @@ -8,4 +8,4 @@ object DummyTrafficAccount : TrafficAccounting { override fun recordFromInternet(bytes: Long) { // do nothing } -} \ No newline at end of file +} diff --git a/core/src/main/kotlin/com/jasonernst/kanonproxy/KAnonProxy.kt b/core/src/main/kotlin/com/jasonernst/kanonproxy/KAnonProxy.kt index 03b7175..d933870 100644 --- a/core/src/main/kotlin/com/jasonernst/kanonproxy/KAnonProxy.kt +++ b/core/src/main/kotlin/com/jasonernst/kanonproxy/KAnonProxy.kt @@ -48,7 +48,7 @@ import java.util.concurrent.atomic.AtomicBoolean class KAnonProxy( val icmp: Icmp, val protector: VpnProtector = DummyProtector, - val trafficAccounting: TrafficAccounting = DummyTrafficAccount + val trafficAccounting: TrafficAccounting = DummyTrafficAccount, ) : SessionManager { private val logger = LoggerFactory.getLogger(javaClass) @@ -213,7 +213,7 @@ class KAnonProxy( protector, this, clientAddress, - trafficAccounting + trafficAccounting, ) } if (isNewSession) { diff --git a/core/src/main/kotlin/com/jasonernst/kanonproxy/TrafficAccounting.kt b/core/src/main/kotlin/com/jasonernst/kanonproxy/TrafficAccounting.kt index dcb57f4..6e2c8ca 100644 --- a/core/src/main/kotlin/com/jasonernst/kanonproxy/TrafficAccounting.kt +++ b/core/src/main/kotlin/com/jasonernst/kanonproxy/TrafficAccounting.kt @@ -2,5 +2,6 @@ package com.jasonernst.kanonproxy interface TrafficAccounting { fun recordToInternet(bytes: Long) + fun recordFromInternet(bytes: Long) -} \ No newline at end of file +} diff --git a/core/src/main/kotlin/com/jasonernst/kanonproxy/tcp/AnonymousTcpSession.kt b/core/src/main/kotlin/com/jasonernst/kanonproxy/tcp/AnonymousTcpSession.kt index dd15202..792ba7f 100644 --- a/core/src/main/kotlin/com/jasonernst/kanonproxy/tcp/AnonymousTcpSession.kt +++ b/core/src/main/kotlin/com/jasonernst/kanonproxy/tcp/AnonymousTcpSession.kt @@ -1,7 +1,6 @@ package com.jasonernst.kanonproxy.tcp import com.jasonernst.kanonproxy.ChangeRequest -import com.jasonernst.kanonproxy.DummyTrafficAccount import com.jasonernst.kanonproxy.SessionManager import com.jasonernst.kanonproxy.TrafficAccounting import com.jasonernst.kanonproxy.VpnProtector @@ -27,7 +26,7 @@ class AnonymousTcpSession( protector: VpnProtector, sessionManager: SessionManager, clientAddress: InetSocketAddress, - trafficAccounting: TrafficAccounting + trafficAccounting: TrafficAccounting, ) : TcpSession( initialIpHeader = initialIpHeader, initialTransportHeader = initialTransportHeader, diff --git a/core/src/main/kotlin/com/jasonernst/kanonproxy/udp/UdpSession.kt b/core/src/main/kotlin/com/jasonernst/kanonproxy/udp/UdpSession.kt index b6bdf39..b455387 100644 --- a/core/src/main/kotlin/com/jasonernst/kanonproxy/udp/UdpSession.kt +++ b/core/src/main/kotlin/com/jasonernst/kanonproxy/udp/UdpSession.kt @@ -31,7 +31,7 @@ class UdpSession( protector: VpnProtector, sessionManager: SessionManager, clientAddress: InetSocketAddress, - trafficAccounting: TrafficAccounting + trafficAccounting: TrafficAccounting, ) : Session( initialIpHeader = initialIpHeader, initialTransportHeader = initialTransportHeader, @@ -40,7 +40,7 @@ class UdpSession( protector = protector, sessionManager = sessionManager, clientAddress = clientAddress, - trafficAccounting = trafficAccounting + trafficAccounting = trafficAccounting, ) { private val logger = LoggerFactory.getLogger(javaClass) diff --git a/core/src/test/kotlin/com/jasonernst/kanonproxy/tcp/AnonymousTcpSessionTest.kt b/core/src/test/kotlin/com/jasonernst/kanonproxy/tcp/AnonymousTcpSessionTest.kt index ae9de27..fcb4b5f 100644 --- a/core/src/test/kotlin/com/jasonernst/kanonproxy/tcp/AnonymousTcpSessionTest.kt +++ b/core/src/test/kotlin/com/jasonernst/kanonproxy/tcp/AnonymousTcpSessionTest.kt @@ -33,7 +33,17 @@ class AnonymousTcpSessionTest { val sessionManager = mockk() every { sessionManager.isRunning() } returns true - val session = AnonymousTcpSession(ipHeader, tcpHeader, ByteArray(0), returnQueue, DummyProtector, sessionManager, clientAddress, DummyTrafficAccount) + val session = + AnonymousTcpSession( + ipHeader, + tcpHeader, + ByteArray(0), + returnQueue, + DummyProtector, + sessionManager, + clientAddress, + DummyTrafficAccount, + ) // wait until its connecting while (session.isConnecting.get().not()) { @@ -54,7 +64,17 @@ class AnonymousTcpSessionTest { ) val tcpHeader2 = TcpHeader(syn = true, destinationPort = 80u) val returnQueue2 = LinkedBlockingDeque() - val session2 = AnonymousTcpSession(ipHeader2, tcpHeader2, ByteArray(0), returnQueue2, DummyProtector, sessionManager, clientAddress, DummyTrafficAccount) + val session2 = + AnonymousTcpSession( + ipHeader2, + tcpHeader2, + ByteArray(0), + returnQueue2, + DummyProtector, + sessionManager, + clientAddress, + DummyTrafficAccount, + ) // wait until its connecting while (session2.isConnecting.get().not()) { diff --git a/core/src/test/kotlin/com/jasonernst/kanonproxy/tcp/TcpClient.kt b/core/src/test/kotlin/com/jasonernst/kanonproxy/tcp/TcpClient.kt index 611945d..906c57d 100644 --- a/core/src/test/kotlin/com/jasonernst/kanonproxy/tcp/TcpClient.kt +++ b/core/src/test/kotlin/com/jasonernst/kanonproxy/tcp/TcpClient.kt @@ -55,7 +55,7 @@ class TcpClient( mockk(relaxed = true), mockk(relaxed = true), clientAddress, - DummyTrafficAccount + DummyTrafficAccount, ) { private val clientId = UUID.randomUUID() private val logger = LoggerFactory.getLogger(javaClass)