Adding a new VPN server

Author: seladb

app/src/main/java/com/pcapplusplus/toyvpn/ToyVpnService.kt

@@ -18,6 +18,7 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.async
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.runBlocking
 import java.io.FileInputStream
 import java.io.FileOutputStream
 import java.io.IOException
@@ -29,7 +30,9 @@ import java.util.concurrent.atomic.AtomicBoolean
 
 class ToyVpnService : VpnService() {
     private var vpnInterface: ParcelFileDescriptor? = null
-    private var vpnConnected: AtomicBoolean = AtomicBoolean(false)
+    private lateinit var vpnTunnel: DatagramChannel
+    private var vpnConnected = AtomicBoolean(false)
+    private var isForwardingTraffic = AtomicBoolean(false)
     private val packetProcessor = PacketProcessor()
     private val packetDataList: MutableList<PacketData> = mutableListOf()
     private var lastPacketDataSentTimestamp: Long = 0
@@ -46,6 +49,7 @@ class ToyVpnService : VpnService() {
         const val MAX_HANDSHAKE_ATTEMPTS = 50
         const val MAX_PACKET_SIZE = 32767
         const val MAX_SECRET_LENGTH = 1024
+        const val DISCONNECT_MESSAGE = "DISCONNECT"
     }
 
     @SuppressLint("SyntheticAccessor")
@@ -71,7 +75,7 @@ class ToyVpnService : VpnService() {
             registerReceiver(
                 broadcastReceiver,
                 IntentFilter(BroadcastActions.VPN_SERVICE_STOP),
-                Context.RECEIVER_EXPORTED,
+                RECEIVER_EXPORTED,
             )
         } else {
             registerReceiver(broadcastReceiver, IntentFilter(BroadcastActions.VPN_SERVICE_STOP))
@@ -99,14 +103,17 @@ class ToyVpnService : VpnService() {
                         }
                     }
 
-                val vpnTunnel = establishVpnResult.await() ?: return@launch
+                vpnTunnel = establishVpnResult.await() ?: return@launch
 
                 try {
+                    isForwardingTraffic.set(true)
                     val inputSteam = FileInputStream(vpnInterface?.fileDescriptor)
                     val outputStream = FileOutputStream(vpnInterface?.fileDescriptor)
                     forwardTraffic(inputSteam, outputStream, vpnTunnel)
                 } catch (e: Exception) {
                     stopSelfOnError("Error while forwarding traffic to the VPN server", e)
+                } finally {
+                    isForwardingTraffic.set(false)
                 }
             }
         } catch (e: IllegalArgumentException) {
@@ -181,6 +188,12 @@ class ToyVpnService : VpnService() {
                     processPacket(packet.array())
                     outputStream.write(packet.array(), 0, length)
                 }
+                else if (length > 1) {
+                    val controlMessage = String(packet.array(), 1, length - 1, US_ASCII)
+                    if (controlMessage == DISCONNECT_MESSAGE) {
+                        stopSelfOnError("Server disconnected")
+                    }
+                }
 
                 idle = false
             }
@@ -289,18 +302,39 @@ class ToyVpnService : VpnService() {
     private fun disconnectVpn() {
         try {
             vpnConnected.set(false)
-            vpnInterface?.close()
-            vpnInterface = null
+
+            runBlocking {
+                vpnServiceScope.async {
+                    while (isForwardingTraffic.get()) {
+                        Thread.sleep(100)
+                    }
+
+                    if (::vpnTunnel.isInitialized) {
+                        val packet = ByteBuffer.allocate(DISCONNECT_MESSAGE.length + 1)
+                        val disconnectAsByteArray = byteArrayOf(0) + DISCONNECT_MESSAGE.encodeToByteArray()
+                        packet.put(disconnectAsByteArray).flip()
+
+                        if (vpnTunnel.write(packet) != disconnectAsByteArray.size) {
+                            Log.e(LOG_TAG, "Couldn't send disconnect message to the server")
+                        }
+                        vpnTunnel.close()
+                        vpnInterface?.close()
+                        vpnInterface = null
+                    }
+                }.await()
+            }
         } catch (e: IOException) {
             e.printStackTrace()
         }
     }
 
-    private fun stopSelfOnError(
+   private fun stopSelfOnError(
         errorMessage: String,
-        exception: java.lang.Exception,
+        exception: java.lang.Exception? = null,
     ) {
-        Log.e(LOG_TAG, errorMessage, exception)
+        if (exception != null) {
+            Log.e(LOG_TAG, errorMessage, exception)
+        }
 
         val intent =
             Intent(BroadcastActions.VPN_SERVICE_ERROR).apply {

server/linux/.gitignore

@@ -0,0 +1,3 @@
+**/*.pcapng
+libs/pcapplusplus/*
+build/*

server/linux/CMakeLists.txt

@@ -0,0 +1,25 @@
+cmake_minimum_required(VERSION 3.20)
+project(ToyVpnServer)
+
+set(CMAKE_CXX_STANDARD 17)
+
+set(PCAPPLUSPLUS_INCLUDE_DIR "${CMAKE_SOURCE_DIR}/libs/PcapPlusPlus/include")
+set(PCAPPLUSPLUS_LIB_DIR "${CMAKE_SOURCE_DIR}/libs/PcapPlusPlus/lib")
+
+# Create the executable target first
+add_executable(ToyVpnServer
+        ToyVpnServer.h
+        main.cpp)
+
+# Add the argparse directory to the target's include path
+target_include_directories(ToyVpnServer PRIVATE argparse)
+
+# Add PcapPlusPlus headers to the include search path
+target_include_directories(ToyVpnServer PRIVATE ${PCAPPLUSPLUS_INCLUDE_DIR})
+
+# Link the necessary PcapPlusPlus libraries
+target_link_libraries(ToyVpnServer PRIVATE
+        ${PCAPPLUSPLUS_LIB_DIR}/libPcap++.a
+        ${PCAPPLUSPLUS_LIB_DIR}/libPacket++.a
+        ${PCAPPLUSPLUS_LIB_DIR}/libCommon++.a
+        pcap)

server/linux/ClientHandler.h

@@ -0,0 +1,148 @@
+#pragma once
+
+#include <chrono>
+#include <netinet/in.h>
+#include "libs/pcapplusplus/include/pcapplusplus/IpAddress.h"
+#include "TunInterfaceWrapper.h"
+#include "ServerSocketWrapper.h"
+#include "VpnSettings.h"
+#include "PacketHandler.h"
+#include "Log.h"
+
+class ClientHandler {
+public:
+    ClientHandler(const ServerSocketWrapper& serverSocket, const sockaddr_in6& clientExternalAddress, const TunInterfaceWrapper& tunInterface, const VpnSettings& vpnSettings, std::optional<PacketHandler>& packetHandler)
+        : m_ServerSocket(serverSocket), m_ClientExternalAddress(clientExternalAddress), m_TunInterface(tunInterface),
+          m_VpnSettings(vpnSettings), m_PacketHandler(packetHandler)
+    {}
+
+    template<std::size_t BUFFER_SIZE>
+    void handleDataFromClient(const std::array<uint8_t, BUFFER_SIZE>& buffer, size_t dataSize) {
+         m_LastMessageTimestamp = std::chrono::steady_clock::now();
+        switch (m_State) {
+            case State::START: {
+                if (dataSize < 2 || buffer[0] != 0) {
+                    break;
+                }
+
+                std::string secret(buffer.begin() + 1, buffer.begin() + dataSize);
+                if (secret != m_VpnSettings.secret) {
+                    TOYVPN_LOG_ERROR("Got the wrong secret: '" << secret << "'");
+                    m_State = State::ERROR;
+                    break;
+                }
+
+                auto params = m_VpnSettings.toParamString();
+                TOYVPN_LOG_DEBUG("Sending params to client: '" << params << "'");
+                std::vector<uint8_t> paramsMessage;
+                paramsMessage.push_back(0);
+                paramsMessage.insert(paramsMessage.end(), params.begin(), params.end());
+                auto bytesSent = m_ServerSocket.send(paramsMessage, m_ClientExternalAddress);
+                if (bytesSent == -1) {
+                    m_State = State::ERROR;
+                } else {
+                    m_State = State::CONNECTED;
+
+                    std::array<uint8_t, 16> ipv6AddressBytes;
+                    std::copy(std::begin(m_ClientExternalAddress.sin6_addr.s6_addr),
+                              std::end(m_ClientExternalAddress.sin6_addr.s6_addr), ipv6AddressBytes.begin());
+
+                    TOYVPN_LOG_INFO("New client connected! External address: ("
+                        << pcpp::IPv6Address(ipv6AddressBytes).toString()
+                        << ","
+                        << m_ClientExternalAddress.sin6_port
+                        << ")"
+                        << ", Internal address: "
+                        << m_VpnSettings.clientAddress);
+                }
+                break;
+            }
+            case State::CONNECTED: {
+                // Got a control packet
+                if (dataSize == 1 && buffer[0] == 0) {
+                    break;
+                }
+
+                if (dataSize == 11 && buffer[0] == 0) {
+                    std::string message(buffer.begin() + 1, buffer.begin() + dataSize);
+                    if (message == m_DisconnectMessage) {
+                        m_State = State::DISCONNECTED;
+
+                        if (m_PacketHandler.has_value()) {
+                            m_PacketHandler->clientDisconnected(m_VpnSettings.clientAddress);
+                        }
+
+                        TOYVPN_LOG_INFO("Client disconnected: " << m_VpnSettings.clientAddress);
+                        break;
+                    }
+                }
+
+                m_TunInterface.send(buffer, dataSize);
+
+                if (m_PacketHandler.has_value()) {
+                    m_PacketHandler->handlePacket(m_VpnSettings.clientAddress, buffer, dataSize);
+                }
+
+                break;
+            }
+            default: {
+            }
+        }
+    }
+
+    template<std::size_t BUFFER_SIZE>
+    void handleDataFromTun(const std::array<uint8_t, BUFFER_SIZE>& buffer, size_t dataSize) {
+        m_ServerSocket.send(buffer, dataSize, m_ClientExternalAddress);
+
+        if (m_PacketHandler.has_value()) {
+            m_PacketHandler->handlePacket(m_VpnSettings.clientAddress, buffer, dataSize);
+        }
+    }
+
+    void disconnect() {
+        if (m_State != State::CONNECTED) {
+            return;
+        }
+
+        std::vector<uint8_t> disconnectMessage;
+        disconnectMessage.push_back(0);
+        disconnectMessage.insert(disconnectMessage.end(), m_DisconnectMessage.begin(), m_DisconnectMessage.end());
+        for (int i = 0; i < 3; i++) {
+            m_ServerSocket.send(disconnectMessage, m_ClientExternalAddress);
+        }
+
+        m_State = State::DISCONNECTED;
+
+        if (m_PacketHandler.has_value()) {
+            m_PacketHandler->clientDisconnected(m_VpnSettings.clientAddress);
+        }
+
+        TOYVPN_LOG_INFO("Client disconnected: " << m_VpnSettings.clientAddress);
+    }
+
+    pcpp::IPv4Address getClientVpnAddress() const { return m_VpnSettings.clientAddress; }
+
+    bool isIdle(const std::chrono::steady_clock::time_point& now) {
+        return m_State == State::DISCONNECTED || now - m_LastMessageTimestamp > m_ClientIdleTimeoutSec;
+    }
+
+private:
+    enum class State {
+        START,
+        CONNECTED,
+        DISCONNECTED,
+        ERROR
+    };
+
+    constexpr static std::string_view m_DisconnectMessage = "DISCONNECT";
+    constexpr static std::chrono::duration m_ClientIdleTimeoutSec = std::chrono::seconds(60);
+
+    const ServerSocketWrapper& m_ServerSocket;
+    const TunInterfaceWrapper& m_TunInterface;
+    std::optional<PacketHandler>& m_PacketHandler;
+    int m_BufferSize;
+    State m_State = State::START;
+    sockaddr_in6 m_ClientExternalAddress;
+    VpnSettings m_VpnSettings;
+    std::chrono::steady_clock::time_point m_LastMessageTimestamp;
+};

server/linux/EpollWrapper.h

@@ -0,0 +1,82 @@
+#pragma once
+
+#include <functional>
+#include <unistd.h>
+#include <sys/epoll.h>
+#include <unordered_map>
+
+class EPollWrapper {
+public:
+    using EPollCallback = std::function<void(int)>;
+
+    virtual ~EPollWrapper() {
+        if (m_EPollFd != -1) {
+            close(m_EPollFd);
+        }
+    }
+
+    void init(int maxEvents) {
+        m_EPollFd = epoll_create1(0);
+        if (m_EPollFd == -1) {
+            throw std::runtime_error("Error creating epoll instance!");
+        }
+
+        m_MaxEvents = maxEvents;
+    }
+
+    void add(int fd, const EPollCallback& callback) {
+        if (m_EPollFd == -1) {
+            throw std::runtime_error("Instance not initialized, please call init()!");
+        }
+
+        struct epoll_event event;
+        event.events = EPOLLIN;
+        event.data.fd = fd;
+        if (epoll_ctl(m_EPollFd, EPOLL_CTL_ADD, fd, &event) == -1) {
+            throw std::runtime_error("Error adding fd to epoll!");
+        }
+
+        m_FdToCallbackMap[event.data.fd] = callback;
+    }
+
+    void remove(int fd) {
+        m_FdToCallbackMap.erase(fd);
+    }
+
+    void startPolling() {
+        if (m_IsPolling) {
+            throw std::runtime_error("Already polling!");
+        }
+
+        m_IsPolling = true;
+
+        epoll_event events[m_MaxEvents];
+
+        while (m_IsPolling) {
+            int numEvents = epoll_wait(m_EPollFd, events, m_MaxEvents, -1);
+            if (numEvents < 0) {
+                if (!m_IsPolling) {
+                    return;
+                }
+                m_IsPolling = false;
+                throw std::runtime_error("Error with epoll_wait!");
+            }
+
+            for (int i = 0; i < numEvents; ++i) {
+                m_FdToCallbackMap[events[i].data.fd](events[i].data.fd);
+            }
+        }
+    }
+
+    void stopPolling() {
+        m_IsPolling = false;
+        close(m_EPollFd);
+    }
+
+private:
+    int m_EPollFd = -1;
+    int m_MaxEvents = -1;
+    std::unordered_map<int, EPollCallback> m_FdToCallbackMap;
+    bool m_IsPolling = false;
+};
+