From 00ba950660d2c182c3f84c7a4b9d423e6ac902b4 Mon Sep 17 00:00:00 2001
From: Connection Refused <cutecr@pm.me>
Date: Mon, 4 Oct 2021 21:36:01 +0800
Subject: [PATCH] [Redirector] Add some functions

---
 Redirector/EventHandler.cpp | 21 +++++++++++++++++----
 Redirector/Redirector.cpp   | 23 +++++++++++++++++++++++
 2 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/Redirector/EventHandler.cpp b/Redirector/EventHandler.cpp
index acfe1706..a342062a 100644
--- a/Redirector/EventHandler.cpp
+++ b/Redirector/EventHandler.cpp
@@ -15,6 +15,9 @@ extern USHORT tcpListen;
 mutex udpContextLock;
 map<ENDPOINT_ID, SocksHelper::PUDP> udpContext;
 
+atomic_ullong UP = { 0 };
+atomic_ullong DL = { 0 };
+
 wstring ConvertIP(PSOCKADDR addr)
 {
 	WCHAR buffer[MAX_PATH] = L"";
@@ -105,6 +108,9 @@ bool eh_init()
 void eh_free()
 {
 	TCPHandler::Free();
+
+	UP = 0;
+	DL = 0;
 }
 
 void threadStart()
@@ -173,7 +179,6 @@ void tcpConnectRequest(ENDPOINT_ID id, PNF_TCP_CONN_INFO info)
 	}
 
 	TCPHandler::CreateHandler(client, remote);
-
 	wcout << "[Redirector][EventHandler][tcpConnectRequest][" << id << "][" << info->processId << "] " << ConvertIP((PSOCKADDR)&client) << " -> " << ConvertIP((PSOCKADDR)&remote) << endl;
 }
 
@@ -189,6 +194,8 @@ void tcpCanSend(ENDPOINT_ID id)
 
 void tcpSend(ENDPOINT_ID id, const char* buffer, int length)
 {
+	UP += length;
+
 	nf_tcpPostSend(id, buffer, length);
 }
 
@@ -199,6 +206,8 @@ void tcpCanReceive(ENDPOINT_ID id)
 
 void tcpReceive(ENDPOINT_ID id, const char* buffer, int length)
 {
+	DL += length;
+
 	nf_tcpPostReceive(id, buffer, length);
 }
 
@@ -251,9 +260,10 @@ void udpSend(ENDPOINT_ID id, const unsigned char* target, const char* buffer, in
 {
 	if (filterDNS && DNSHandler::IsDNS((PSOCKADDR_IN6)target))
 	{
-		wcout << "[Redirector][EventHandler][udpSend][" << id << "] DNS to " << ConvertIP((PSOCKADDR)target) << endl;
-
+		UP += length;
 		DNSHandler::CreateHandler(id, (PSOCKADDR_IN6)target, buffer, length, options);
+
+		wcout << "[Redirector][EventHandler][udpSend][" << id << "] DNS to " << ConvertIP((PSOCKADDR)target) << endl;
 		return;
 	}
 
@@ -265,10 +275,11 @@ void udpSend(ENDPOINT_ID id, const unsigned char* target, const char* buffer, in
 		nf_udpPostSend(id, target, buffer, length, options);
 		return;
 	}
-
 	auto udpConn = udpContext[id];
 	udpContextLock.unlock();
 
+	UP += length;
+
 	if (udpConn->tcpSocket == INVALID_SOCKET)
 	{
 		auto tcpSocket = SocksHelper::Utils::Connect();
@@ -366,6 +377,8 @@ void udpBeginReceive(ENDPOINT_ID id, SocksHelper::PUDP udpConn, PNF_UDP_OPTIONS
 			break;
 		}
 
+		DL += length;
+
 		nf_udpPostReceive(id, (unsigned char*)&target, buffer, length, options);
 	}
 
diff --git a/Redirector/Redirector.cpp b/Redirector/Redirector.cpp
index fbdc406a..a53389c4 100644
--- a/Redirector/Redirector.cpp
+++ b/Redirector/Redirector.cpp
@@ -23,6 +23,9 @@ extern string tgtPassword;
 extern vector<wstring> bypassList;
 extern vector<wstring> handleList;
 
+extern atomic_ullong UP;
+extern atomic_ullong DL;
+
 NF_EventHandler EventHandler = {
 	threadStart,
 	threadEnd,
@@ -57,6 +60,16 @@ BOOL APIENTRY DllMain(HMODULE hModule, DWORD dwReason, LPVOID lpReserved)
 }
 
 extern "C" {
+	__declspec(dllexport) BOOL __cdecl aio_register(LPWSTR value)
+	{
+		return nf_registerDriver(ws2s(value).c_str()) == NF_STATUS_SUCCESS;
+	}
+
+	__declspec(dllexport) BOOL __cdecl aio_unregister(LPWSTR value)
+	{
+		return nf_unRegisterDriver(ws2s(value).c_str()) == NF_STATUS_SUCCESS;
+	}
+
 	__declspec(dllexport) BOOL __cdecl aio_dial(int name, LPWSTR value)
 	{
 		switch (name)
@@ -288,4 +301,14 @@ extern "C" {
 		WSACleanup();
 		return;
 	}
+
+	__declspec(dllexport) ULONG64 __cdecl aio_getUP()
+	{
+		return UP;
+	}
+
+	__declspec(dllexport) ULONG64 __cdecl aio_getDL()
+	{
+		return DL;
+	}
 }