From fabf6f51cdcaa6e4c2face36c1ea57b318d63bab Mon Sep 17 00:00:00 2001
From: Sepalani <sepalani@hotmail.fr>
Date: Thu, 11 Feb 2021 16:13:41 +0400
Subject: [PATCH] PCAP: Improve TCP session logging

---
 Source/Core/Core/IOS/Network/Socket.cpp   |  1 +
 Source/Core/Core/NetworkCaptureLogger.cpp | 14 ++++++++++++--
 Source/Core/Core/NetworkCaptureLogger.h   | 11 +++++++++--
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/Source/Core/Core/IOS/Network/Socket.cpp b/Source/Core/Core/IOS/Network/Socket.cpp
index 3cf193cb57..1c01bdabeb 100644
--- a/Source/Core/Core/IOS/Network/Socket.cpp
+++ b/Source/Core/Core/IOS/Network/Socket.cpp
@@ -701,6 +701,7 @@ s32 WiiSockMan::AddSocket(s32 fd, bool is_rw)
     WiiSocket& sock = WiiSockets[wii_fd];
     sock.SetFd(fd);
     sock.SetWiiFd(wii_fd);
+    PowerPC::debug_interface.NetworkLogger()->OnNewSocket(fd);
 
 #ifdef __APPLE__
     int opt_no_sigpipe = 1;
diff --git a/Source/Core/Core/NetworkCaptureLogger.cpp b/Source/Core/Core/NetworkCaptureLogger.cpp
index 048adb4e08..118deb234a 100644
--- a/Source/Core/Core/NetworkCaptureLogger.cpp
+++ b/Source/Core/Core/NetworkCaptureLogger.cpp
@@ -24,6 +24,10 @@ namespace Core
 NetworkCaptureLogger::NetworkCaptureLogger() = default;
 NetworkCaptureLogger::~NetworkCaptureLogger() = default;
 
+void DummyNetworkCaptureLogger::OnNewSocket(s32 socket)
+{
+}
+
 void DummyNetworkCaptureLogger::LogSSLRead(const void* data, std::size_t length, s32 socket)
 {
 }
@@ -81,6 +85,12 @@ PCAPSSLCaptureLogger::PCAPSSLCaptureLogger()
 
 PCAPSSLCaptureLogger::~PCAPSSLCaptureLogger() = default;
 
+void PCAPSSLCaptureLogger::OnNewSocket(s32 socket)
+{
+  m_read_sequence_number[socket] = 0;
+  m_write_sequence_number[socket] = 0;
+}
+
 PCAPSSLCaptureLogger::ErrorState PCAPSSLCaptureLogger::SaveState() const
 {
   return {
@@ -193,8 +203,8 @@ void PCAPSSLCaptureLogger::LogIPv4(LogType log_type, const u8* data, u16 length,
 
   if (socket_type == SOCK_STREAM)
   {
-    u32& sequence_number =
-        (log_type == LogType::Read) ? m_read_sequence_number : m_write_sequence_number;
+    u32& sequence_number = (log_type == LogType::Read) ? m_read_sequence_number[socket] :
+                                                         m_write_sequence_number[socket];
     Common::TCPHeader tcp_header(from, to, sequence_number, data, length);
     sequence_number += static_cast<u32>(length);
     Common::IPv4Header ip_header(tcp_header.Size() + length, tcp_header.IPProto(), from, to);
diff --git a/Source/Core/Core/NetworkCaptureLogger.h b/Source/Core/Core/NetworkCaptureLogger.h
index 814c630855..9dbca686ca 100644
--- a/Source/Core/Core/NetworkCaptureLogger.h
+++ b/Source/Core/Core/NetworkCaptureLogger.h
@@ -6,6 +6,7 @@
 
 #include <array>
 #include <cstddef>
+#include <map>
 #include <memory>
 
 #ifdef _WIN32
@@ -43,6 +44,8 @@ public:
   NetworkCaptureLogger& operator=(NetworkCaptureLogger&&) = delete;
   virtual ~NetworkCaptureLogger();
 
+  virtual void OnNewSocket(s32 socket) = 0;
+
   virtual void LogSSLRead(const void* data, std::size_t length, s32 socket) = 0;
   virtual void LogSSLWrite(const void* data, std::size_t length, s32 socket) = 0;
 
@@ -55,6 +58,8 @@ public:
 class DummyNetworkCaptureLogger : public NetworkCaptureLogger
 {
 public:
+  void OnNewSocket(s32 socket) override;
+
   void LogSSLRead(const void* data, std::size_t length, s32 socket) override;
   void LogSSLWrite(const void* data, std::size_t length, s32 socket) override;
 
@@ -79,6 +84,8 @@ public:
   PCAPSSLCaptureLogger();
   ~PCAPSSLCaptureLogger();
 
+  void OnNewSocket(s32 socket) override;
+
   void LogSSLRead(const void* data, std::size_t length, s32 socket) override;
   void LogSSLWrite(const void* data, std::size_t length, s32 socket) override;
 
@@ -108,7 +115,7 @@ private:
                const sockaddr_in& to);
 
   std::unique_ptr<Common::PCAP> m_file;
-  u32 m_read_sequence_number = 0;
-  u32 m_write_sequence_number = 0;
+  std::map<s32, u32> m_read_sequence_number;
+  std::map<s32, u32> m_write_sequence_number;
 };
 }  // namespace Core