Skip to content

Commit

Permalink
fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
tebaikin committed Nov 27, 2023
1 parent c4b1476 commit 8779510
Showing 1 changed file with 117 additions and 140 deletions.
257 changes: 117 additions & 140 deletions Vostok.Metrics.System/Host/TcpStateCollector.cs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
using System;
using System.Collections.Generic;
using System.Net;
using System.Net.NetworkInformation;
using System.Net.Sockets;
using System.Runtime.InteropServices;
Expand All @@ -13,44 +12,48 @@ public void Collect(HostMetrics metrics)
{
var states = new Dictionary<TcpState, int>();

#if NET6_0_OR_GREATER
CountStates(states);

metrics.TcpStates = states;
}

private static void CountStates(IDictionary<TcpState, int> states)
{
#if NET6_0_OR_GREATER

if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
{
IterateOverTcpConnections(state =>
{
if (!states.ContainsKey(state))
states[state] = 0;
states[state]++;
if (state is not TcpState.Listen)
CountState(state);
});
}
else
{
foreach (var tcpConnection in IPGlobalProperties.GetIPGlobalProperties().GetActiveTcpConnections())
{
if (!states.ContainsKey(tcpConnection.State))
states[tcpConnection.State] = 0;
states[tcpConnection.State]++;
CountState(tcpConnection.State);
}
}

#else

#else
foreach (var tcpConnection in IPGlobalProperties.GetIPGlobalProperties().GetActiveTcpConnections())
{
if (!states.ContainsKey(tcpConnection.State))
states[tcpConnection.State] = 0;
states[tcpConnection.State]++;
CountState(tcpConnection.State);
}

#endif



metrics.TcpStates = states;
#endif

void CountState(TcpState state)
{
if (!states.ContainsKey(state))
states[state] = 0;
states[state]++;
}
}
#if NET6_0_OR_GREATER

#if NET6_0_OR_GREATER
private static unsafe void IterateOverTcpConnections(Action<TcpState> handle)
{
var size = 0U;
Expand All @@ -71,28 +74,16 @@ private static unsafe void IterateOverTcpConnections(Action<TcpState> handle)
{
var span = new ReadOnlySpan<byte>((byte*)buffer, (int)size);

// The table info just gives us the number of rows.
ref readonly var tcpTableInfo = ref MemoryMarshal.AsRef<Interop.IpHlpApi.MibTcpTable>(span);

if (tcpTableInfo.numberOfEntries > 0)
{
var row = MemoryMarshal.AsRef<Interop.IpHlpApi.MibTcpRow>(span);
// Skip over the tableinfo to get the inline rows.
span = span[sizeof(Interop.IpHlpApi.MibTcpTable)..];

for (var i = 0; i < tcpTableInfo.numberOfEntries; i++)
{
var state = row.state;

// Port is returned in Big-Endian - most significant bit on left.
// Unfortunately, its done at the word level and not the DWORD level.
var localPort = row.localPort1 << 8 | row.localPort2;
var remotePort = state == TcpState.Listen ? 0 : row.remotePort1 << 8 | row.remotePort2;

// var info = new TcpInfo(state, new IPEndPoint(row.localAddr, localPort), new IPEndPoint(row.remoteAddr, remotePort));

handle(state);

var row = MemoryMarshal.AsRef<Interop.IpHlpApi.MibTcpRow>(span);
handle(row.state);
span = span[sizeof(Interop.IpHlpApi.MibTcpRow)..];
}
}
Expand All @@ -104,7 +95,6 @@ private static unsafe void IterateOverTcpConnections(Action<TcpState> handle)
}
}

// If we don't have any ipv4 interfaces detected, just continue.
if (result != Interop.IpHlpApi.ERROR_SUCCESS && result != Interop.IpHlpApi.ERROR_NO_DATA)
{
throw new NetworkInformationException((int)result);
Expand All @@ -113,16 +103,14 @@ private static unsafe void IterateOverTcpConnections(Action<TcpState> handle)

if (Socket.OSSupportsIPv6)
{
// Get the buffer size needed.
size = 0;
result = Interop.IpHlpApi.GetExtendedTcpTable(IntPtr.Zero, ref size, true,
(uint)AddressFamily.InterNetworkV6,
Interop.IpHlpApi.TcpTableClass.TcpTableOwnerPidAll, 0);

while (result == Interop.IpHlpApi.ERROR_INSUFFICIENT_BUFFER)
{
// Allocate the buffer and get the TCP table.
IntPtr buffer = Marshal.AllocHGlobal((int)size);
var buffer = Marshal.AllocHGlobal((int)size);
try
{
result = Interop.IpHlpApi.GetExtendedTcpTable(buffer, ref size, true,
Expand All @@ -132,32 +120,17 @@ private static unsafe void IterateOverTcpConnections(Action<TcpState> handle)
{
var span = new ReadOnlySpan<byte>((byte*)buffer, (int)size);

// The table info just gives us the number of rows.
ref readonly Interop.IpHlpApi.MibTcp6TableOwnerPid tcpTable6OwnerPid = ref MemoryMarshal.AsRef<Interop.IpHlpApi.MibTcp6TableOwnerPid>(span);
ref readonly var tcpTable6OwnerPid = ref MemoryMarshal.AsRef<Interop.IpHlpApi.MibTcp6TableOwnerPid>(span);

if (tcpTable6OwnerPid.numberOfEntries > 0)
{
// Skip over the tableinfo to get the inline rows.
span = span.Slice(sizeof(Interop.IpHlpApi.MibTcp6TableOwnerPid));
span = span[sizeof(Interop.IpHlpApi.MibTcp6TableOwnerPid)..];

for (var i = 0; i < tcpTable6OwnerPid.numberOfEntries; i++)
{
var row = MemoryMarshal.AsRef<Interop.IpHlpApi.MibTcp6RowOwnerPid>(span);
var state = row.state;

// Port is returned in Big-Endian - most significant bit on left.
// Unfortunately, its done at the word level and not the DWORD level.
var localPort = row.localPort1 << 8 | row.localPort2;
var remotePort = ((state == TcpState.Listen) ? 0 : row.remotePort1 << 8 | row.remotePort2);

// var info = new TcpInfo(
// state,
// new IPEndPoint(new IPAddress(row.localAddrAsSpan, row.localScopeId), localPort),
// new IPEndPoint(new IPAddress(row.remoteAddrAsSpan, row.remoteScopeId), remotePort));

handle(state);
// We increment the pointer to the next row.
span = span.Slice(sizeof(Interop.IpHlpApi.MibTcp6RowOwnerPid));
handle(row.state);
span = span[sizeof(Interop.IpHlpApi.MibTcp6RowOwnerPid)..];
}
}
}
Expand All @@ -168,7 +141,6 @@ private static unsafe void IterateOverTcpConnections(Action<TcpState> handle)
}
}

// If we don't have any ipv6 interfaces detected, just continue.
if (result != Interop.IpHlpApi.ERROR_SUCCESS && result != Interop.IpHlpApi.ERROR_NO_DATA)
{
throw new NetworkInformationException((int)result);
Expand All @@ -177,99 +149,104 @@ private static unsafe void IterateOverTcpConnections(Action<TcpState> handle)
}
#endif
}

#if NET6_0_OR_GREATER
internal static class Interop
{
private class Libraries
{
public const string IpHlpApi = "iphlpapi.dll";
}

internal static class IpHlpApi

#if NET6_0_OR_GREATER
internal static class Interop
{
public const int ERROR_INSUFFICIENT_BUFFER = 0x007A;
public const int ERROR_SUCCESS = 0x0;
public const int ERROR_NO_DATA = 0xE8;

[StructLayout(LayoutKind.Sequential)]
internal struct MibTcpTable
private class Libraries
{
internal uint numberOfEntries;
public const string IpHlpApi = "iphlpapi.dll";
}

[StructLayout(LayoutKind.Sequential)]
internal struct MibTcpRow
internal static class IpHlpApi
{
internal TcpState state;
internal uint localAddr;
internal byte localPort1;
internal byte localPort2;
// Ports are only 16 bit values (in network WORD order, 3,4,1,2).
// There are reports where the high order bytes have garbage in them.
internal byte ignoreLocalPort3;
internal byte ignoreLocalPort4;
internal uint remoteAddr;
internal byte remotePort1;
internal byte remotePort2;
// Ports are only 16 bit values (in network WORD order, 3,4,1,2).
// There are reports where the high order bytes have garbage in them.
internal byte ignoreRemotePort3;
internal byte ignoreRemotePort4;
}
public const int ERROR_INSUFFICIENT_BUFFER = 0x007A;
public const int ERROR_SUCCESS = 0x0;
public const int ERROR_NO_DATA = 0xE8;

[StructLayout(LayoutKind.Sequential)]
internal struct MibTcp6TableOwnerPid
{
internal uint numberOfEntries;
}
[StructLayout(LayoutKind.Sequential)]
internal struct MibTcpTable
{
internal uint numberOfEntries;
}

[StructLayout(LayoutKind.Sequential)]
internal unsafe struct MibTcp6RowOwnerPid
{
internal fixed byte localAddr[16];
internal uint localScopeId;
internal byte localPort1;
internal byte localPort2;
// Ports are only 16 bit values (in network WORD order, 3,4,1,2).
// There are reports where the high order bytes have garbage in them.
internal byte ignoreLocalPort3;
internal byte ignoreLocalPort4;
internal fixed byte remoteAddr[16];
internal uint remoteScopeId;
internal byte remotePort1;
internal byte remotePort2;
// Ports are only 16 bit values (in network WORD order, 3,4,1,2).
// There are reports where the high order bytes have garbage in them.
internal byte ignoreRemotePort3;
internal byte ignoreRemotePort4;
internal TcpState state;
internal uint owningPid;
[StructLayout(LayoutKind.Sequential)]
internal struct MibTcpRow
{
internal TcpState state;
internal uint localAddr;
internal byte localPort1;
internal byte localPort2;
// Ports are only 16 bit values (in network WORD order, 3,4,1,2).
// There are reports where the high order bytes have garbage in them.
internal byte ignoreLocalPort3;
internal byte ignoreLocalPort4;
internal uint remoteAddr;
internal byte remotePort1;
internal byte remotePort2;
// Ports are only 16 bit values (in network WORD order, 3,4,1,2).
// There are reports where the high order bytes have garbage in them.
internal byte ignoreRemotePort3;
internal byte ignoreRemotePort4;
}

internal ReadOnlySpan<byte> localAddrAsSpan => MemoryMarshal.CreateSpan(ref localAddr[0], 16);
internal ReadOnlySpan<byte> remoteAddrAsSpan => MemoryMarshal.CreateSpan(ref remoteAddr[0], 16);
}
[StructLayout(LayoutKind.Sequential)]
internal struct MibTcp6TableOwnerPid
{
internal uint numberOfEntries;
}

internal enum TcpTableClass
{
TcpTableBasicListener = 0,
TcpTableBasicConnections = 1,
TcpTableBasicAll = 2,
TcpTableOwnerPidListener = 3,
TcpTableOwnerPidConnections = 4,
TcpTableOwnerPidAll = 5,
TcpTableOwnerModuleListener = 6,
TcpTableOwnerModuleConnections = 7,
TcpTableOwnerModuleAll = 8
}
[StructLayout(LayoutKind.Sequential)]
internal unsafe struct MibTcp6RowOwnerPid
{
internal fixed byte localAddr[16];
internal uint localScopeId;
internal byte localPort1;
internal byte localPort2;
// Ports are only 16 bit values (in network WORD order, 3,4,1,2).
// There are reports where the high order bytes have garbage in them.
internal byte ignoreLocalPort3;
internal byte ignoreLocalPort4;
internal fixed byte remoteAddr[16];
internal uint remoteScopeId;
internal byte remotePort1;
internal byte remotePort2;
// Ports are only 16 bit values (in network WORD order, 3,4,1,2).
// There are reports where the high order bytes have garbage in them.
internal byte ignoreRemotePort3;
internal byte ignoreRemotePort4;
internal TcpState state;
internal uint owningPid;

internal ReadOnlySpan<byte> localAddrAsSpan => MemoryMarshal.CreateSpan(ref localAddr[0], 16);
internal ReadOnlySpan<byte> remoteAddrAsSpan => MemoryMarshal.CreateSpan(ref remoteAddr[0], 16);
}

[DllImport(Libraries.IpHlpApi)]
internal static extern uint GetTcpTable(IntPtr pTcpTable, ref uint dwOutBufLen, bool order);
internal enum TcpTableClass
{
TcpTableBasicListener = 0,
TcpTableBasicConnections = 1,
TcpTableBasicAll = 2,
TcpTableOwnerPidListener = 3,
TcpTableOwnerPidConnections = 4,
TcpTableOwnerPidAll = 5,
TcpTableOwnerModuleListener = 6,
TcpTableOwnerModuleConnections = 7,
TcpTableOwnerModuleAll = 8
}

[DllImport(Libraries.IpHlpApi)]
internal static extern uint GetExtendedTcpTable(IntPtr pTcpTable, ref uint dwOutBufLen, bool order,
uint IPVersion, TcpTableClass tableClass, uint reserved);
[DllImport(Libraries.IpHlpApi)]
internal static extern uint GetTcpTable(IntPtr pTcpTable, ref uint dwOutBufLen, bool order);

[DllImport(Libraries.IpHlpApi)]
internal static extern uint GetExtendedTcpTable(
IntPtr pTcpTable,
ref uint dwOutBufLen,
bool order,
uint IPVersion,
TcpTableClass tableClass,
uint reserved);
}
}
}
#endif
}

0 comments on commit 8779510

Please sign in to comment.