~ruther/NosSmooth

ab84d2f44498c9563c164ed2ebe9bfe941b569d9 — Rutherther 2 years ago b1ca4f8
feat(pcap): add tcp connection to pid helpers
A Pcap/NosSmooth.Pcap/ProcessTcpManager.cs => Pcap/NosSmooth.Pcap/ProcessTcpManager.cs +113 -0
@@ 0,0 1,113 @@
//
//  ProcessTcpManager.cs
//
//  Copyright (c) František Boháček. All rights reserved.
//  Licensed under the MIT license. See LICENSE file in the project root for full license information.

using System.Diagnostics;

namespace NosSmooth.Pcap;

/// <summary>
/// A manager containing tcp connections, allowing notifications
/// to <see cref="PcapNostaleClient"/> to know about any new connections.
/// </summary>
public class ProcessTcpManager
{
    private static TimeSpan RefreshInterval = TimeSpan.FromMilliseconds(0.99);

    private readonly SemaphoreSlim _semaphore;
    private readonly List<int> _processes;
    private DateTimeOffset _lastRefresh;
    private IReadOnlyDictionary<int, List<TcpConnection>> _connections;

    /// <summary>
    /// Initializes a new instance of the <see cref="ProcessTcpManager"/> class.
    /// </summary>
    public ProcessTcpManager()
    {
        _lastRefresh = DateTimeOffset.MinValue;
        _semaphore = new SemaphoreSlim(1, 1);
        _processes = new List<int>();
        _connections = new Dictionary<int, List<TcpConnection>>();
    }

    /// <summary>
    /// Register the given process to refreshing list to allow calling <see cref="GetConnectionsAsync"/>
    /// with that process.
    /// </summary>
    /// <param name="processId">The id of the process to register.</param>
    /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
    public async Task RegisterProcess(int processId)
    {
        await _semaphore.WaitAsync();
        try
        {
            _processes.Add(processId);
        }
        finally
        {
            _semaphore.Release();
        }
    }

    /// <summary>
    /// Unregister the given process from refreshing list, <see cref="GetConnectionsAsync"/> won't
    /// work for that process anymore.
    /// </summary>
    /// <param name="processId">The process to unregister.</param>
    /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
    public async Task UnregisterProcess(int processId)
    {
        await _semaphore.WaitAsync();
        try
        {
            _processes.Remove(processId);
        }
        finally
        {
            _semaphore.Release();
        }
    }

    /// <summary>
    /// Get connections established by the given process.
    /// </summary>
    /// <remarks>
    /// Works only for processes registered using <see cref="RegisterProcess"/>.
    /// </remarks>
    /// <param name="processId">The id of process to retrieve connections for.</param>
    /// <returns>The list of process connections.</returns>
    public async Task<IReadOnlyList<TcpConnection>> GetConnectionsAsync(int processId)
    {
        await Refresh();

        if (!_connections.ContainsKey(processId))
        {
            return Array.Empty<TcpConnection>();
        }

        return _connections[processId];
    }

    private async Task Refresh()
    {
        if (_lastRefresh.Add(RefreshInterval) >= DateTimeOffset.Now)
        {
            return;
        }

        _lastRefresh = DateTimeOffset.Now;
        if (_processes.Count == 0)
        {
            if (_connections.Count > 0)
            {
                _connections = new Dictionary<int, List<TcpConnection>>();
            }
        }

        await _semaphore.WaitAsync();
        _connections = TcpConnectionHelper.GetConnections(_processes);
        _semaphore.Release();
    }
}
\ No newline at end of file

A Pcap/NosSmooth.Pcap/TcpConnection.cs => Pcap/NosSmooth.Pcap/TcpConnection.cs +19 -0
@@ 0,0 1,19 @@
//
//  TcpConnection.cs
//
//  Copyright (c) František Boháček. All rights reserved.
//  Licensed under the MIT license. See LICENSE file in the project root for full license information.

using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;

namespace NosSmooth.Pcap;

[SuppressMessage("StyleCop.CSharp.NamingRules", "SA1313:Parameter names should begin with lower-case letter", Justification = "Fix this.")]
public record struct TcpConnection
(
    long LocalAddr,
    int LocalPort,
    long RemoteAddr,
    int RemotePort
);
\ No newline at end of file

A Pcap/NosSmooth.Pcap/TcpConnectionHelper.cs => Pcap/NosSmooth.Pcap/TcpConnectionHelper.cs +207 -0
@@ 0,0 1,207 @@
//
//  TcpConnectionHelper.cs
//
//  Copyright (c) František Boháček. All rights reserved.
//  Licensed under the MIT license. See LICENSE file in the project root for full license information.

using System.Diagnostics;
using System.Net.NetworkInformation;
using System.Runtime.InteropServices;

namespace NosSmooth.Pcap;

/// <summary>
/// A class for obtaining process tcp connections.
/// </summary>
/// <remarks>
/// Works on Windows only so far.
/// </remarks>
public static class TcpConnectionHelper
{
    private const int AF_INET = 2; // IP_v4 = System.Net.Sockets.AddressFamily.InterNetwork
    private const int AF_INET6 = 23; // IP_v6 = System.Net.Sockets.AddressFamily.InterNetworkV6

    /// <summary>
    /// Get TCP IPv4 connections of the specified processes.
    /// </summary>
    /// <param name="processIds">The process ids to look for.</param>
    /// <returns>Map from process ids to connecitons.</returns>
    /// <exception cref="NotImplementedException">Thrown if not windows.</exception>
    public static IReadOnlyDictionary<int, List<TcpConnection>> GetConnections(IReadOnlyList<int> processIds)
    {
        if (!RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
        {
            throw new NotImplementedException();
        }

        var result = new Dictionary<int, List<TcpConnection>>();
        var tcpv4Connections = GetAllTCPv4Connections();

        foreach (var connection in tcpv4Connections)
        {
            var process = processIds.FirstOrDefault(x => x == connection.OwningPid, -1);
            if (process != -1)
            {
                if (!result.ContainsKey(process))
                {
                    result.Add(process, new List<TcpConnection>());
                }

                result[process].Add
                (
                    new TcpConnection
                    (
                        connection.LocalAddr,
                        (ushort)(connection.LocalPort[1] | (connection.LocalPort[0] << 8)),
                        connection.RemoteAddr,
                        (ushort)(connection.RemotePort[1] | (connection.RemotePort[0] << 8))
                    )
                );
            }
        }

        return result;
    }

    private static List<MIB_TCPROW_OWNER_PID> GetAllTCPv4Connections()
    {
        return GetTCPConnections<MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID>(AF_INET);
    }

    private static List<MIB_TCP6ROW_OWNER_PID> GetAllTCPv6Connections()
    {
        return GetTCPConnections<MIB_TCP6ROW_OWNER_PID, MIB_TCP6TABLE_OWNER_PID>(AF_INET6);
    }

    private static List<TIPR> GetTCPConnections<TIPR, TIPT>(int ipVersion)
    {
        // IPR = Row Type, IPT = Table Type

        TIPR[] tableRows;
        int buffSize = 0;
        var dwNumEntriesField = typeof(TIPT).GetField("DwNumEntries");

        // how much memory do we need?
        uint ret = GetExtendedTcpTable
        (
            IntPtr.Zero,
            ref buffSize,
            true,
            ipVersion,
            TCP_TABLE_CLASS.TCP_TABLE_OWNER_PID_ALL
        );
        IntPtr tcpTablePtr = Marshal.AllocHGlobal(buffSize);

        try
        {
            ret = GetExtendedTcpTable
            (
                tcpTablePtr,
                ref buffSize,
                true,
                ipVersion,
                TCP_TABLE_CLASS.TCP_TABLE_OWNER_PID_ALL
            );
            if (ret != 0)
            {
                return new List<TIPR>();
            }

            // get the number of entries in the table
            TIPT table = (TIPT)Marshal.PtrToStructure(tcpTablePtr, typeof(TIPT))!;
            int rowStructSize = Marshal.SizeOf(typeof(TIPR));
            uint numEntries = (uint)dwNumEntriesField!.GetValue(table)!;

            // buffer we will be returning
            tableRows = new TIPR[numEntries];

            IntPtr rowPtr = (IntPtr)((long)tcpTablePtr + 4);
            for (int i = 0; i < numEntries; i++)
            {
                TIPR tcpRow = (TIPR)Marshal.PtrToStructure(rowPtr, typeof(TIPR))!;
                tableRows[i] = tcpRow;
                rowPtr = (IntPtr)((long)rowPtr + rowStructSize); // next entry
            }
        }
        finally
        {
            // Free the Memory
            Marshal.FreeHGlobal(tcpTablePtr);
        }
        return tableRows != null ? tableRows.ToList() : new List<TIPR>();
    }

    // https://msdn2.microsoft.com/en-us/library/aa366913.aspx
    [StructLayout(LayoutKind.Sequential)]
    private struct MIB_TCPROW_OWNER_PID
    {
        public uint State;
        public uint LocalAddr;
        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 4)]
        public byte[] LocalPort;
        public uint RemoteAddr;
        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 4)]
        public byte[] RemotePort;
        public uint OwningPid;
    }

    // https://msdn2.microsoft.com/en-us/library/aa366921.aspx
    [StructLayout(LayoutKind.Sequential)]
    private struct MIB_TCPTABLE_OWNER_PID
    {
        public uint DwNumEntries;
        [MarshalAs(UnmanagedType.ByValArray, ArraySubType = UnmanagedType.Struct, SizeConst = 1)]
        public MIB_TCPROW_OWNER_PID[] Table;
    }

    // https://msdn.microsoft.com/en-us/library/aa366896
    [StructLayout(LayoutKind.Sequential)]
    private struct MIB_TCP6ROW_OWNER_PID
    {
        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 16)]
        public byte[] LocalAddr;
        public uint LocalScopeId;
        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 4)]
        public byte[] LocalPort;
        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 16)]
        public byte[] RemoteAddr;
        public uint RemoteScopeId;
        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 4)]
        public byte[] RemotePort;
        public uint State;
        public uint OwningPid;
    }

    // https://msdn.microsoft.com/en-us/library/windows/desktop/aa366905
    [StructLayout(LayoutKind.Sequential)]
    private struct MIB_TCP6TABLE_OWNER_PID
    {
        public uint DwNumEntries;
        [MarshalAs(UnmanagedType.ByValArray, ArraySubType = UnmanagedType.Struct, SizeConst = 1)]
        public MIB_TCP6ROW_OWNER_PID[] Table;
    }

    [DllImport("iphlpapi.dll", SetLastError = true)]
    private static extern uint GetExtendedTcpTable
    (
        IntPtr pTcpTable,
        ref int dwOutBufLen,
        bool sort,
        int ipVersion,
        TCP_TABLE_CLASS tblClass,
        uint reserved = 0
    );

    private enum TCP_TABLE_CLASS
    {
        TCP_TABLE_BASIC_LISTENER,
        TCP_TABLE_BASIC_CONNECTIONS,
        TCP_TABLE_BASIC_ALL,
        TCP_TABLE_OWNER_PID_LISTENER,
        TCP_TABLE_OWNER_PID_CONNECTIONS,
        TCP_TABLE_OWNER_PID_ALL,
        TCP_TABLE_OWNER_MODULE_LISTENER,
        TCP_TABLE_OWNER_MODULE_CONNECTIONS,
        TCP_TABLE_OWNER_MODULE_ALL
    }
}
\ No newline at end of file

Do not follow this link