UDP Latency Check

example-server.ps1

#!/usr/bin/env pwsh

# Save results to a psql database
function Save-Result {
    [CmdletBinding()]
    param (
        [Parameter(ValueFromPipeline=$true)]
        [string]$json
    )
    Process {
        $sqlStatement = $json | jq --raw-output @"
{ "src": ( .address | capture("(?<ip>.*):(?<port>.*?$)") ) } +
  (.message | capture("(?<date_reported>[0-9/]+ [0-9:]{8}): ?(?<message>.*?)[ms]{0,2}$"; ""))
  | "INSERT INTO pings (ip, port, message, date_reported) VALUES ('\(.src.ip)', \(.src.port), '\(.message)', '\(.date_reported)');"
"@
        Write-Host $sqlStatement
        psql --dbname stats -c "$sqlStatement"
    }
}

./latency-check.ps1 `
  -Server `
  -AesKey ([System.Convert]::FromBase64String("REDACTED_B64_AES_KEY")) `
  -ListenPort 11001 `
  -ReceiveTimeout 3000 | Save-Result

latency-check.ps1

#!/usr/bin/env pwsh

Param (
    # Port number to listen on
    [int] $ListenPort = 11000,
    # Port number to send data to
    [int] $SendPort = 11000,
    # Timeout in milliseconds for sending data
    [int] $SendTimeout = 1000,
    # Timeout in milliseconds for receiving data
    [int] $ReceiveTimeout = 1000,
    # Address to send data to
    [string] $SendAddress = "127.0.0.1",
    # Message to send
    [string] $Message = "Hello",
    # Base64-encoded message to send
    [string] $MessageBase64 = "",
    [string] $HmacKey = "changeit",
    [byte[]] $AesKey,
    [Alias("Server")]
    [switch] $RunServer
)

$AES = [System.Security.Cryptography.Aes]::Create()
if ($null -ne $aesKey)
{
    $AES.Key = $aesKey
    $AES.GenerateIV()
}

$encoder = [system.Text.Encoding]::UTF8

function Test-Latency {
    param (
        [System.Net.Sockets.UdpClient] $client
    )

    if ( $MessageBase64 -ne "" ) 
    {
        $sendBytes = [Convert]::FromBase64String($MessageBase64)
    }
    else 
    {
        Write-Debug $ListenPort
        # use ascii message
        $sendBytes = $encoder.GetBytes($Message)
    }
    
    if ($null -ne $aesKey) 
    {
        $sendBytes = ( $AES.IV + $AES.EncryptCbc($sendBytes, $AES.IV) )
    }
    
    # Measure the time it takes to send and receive data
    Measure-Command {
        # Send the message
        $transmittedBytes = $client.Send($sendBytes, $sendBytes.Length, $SendAddress, $SendPort)
    
        # Check if the message was sent successfully
        if ($transmittedBytes -gt 0)
        {
            $RemoteIpEndPoint = [System.Net.IPEndPoint]::new([System.Net.IPAddress]::Any, 0)
        
            try
            {
                # Receive the response message
                $receiveBytes = $client.Receive([ref]$RemoteIpEndPoint);
    
                # Convert the response message to a string
                $returnData = $encoder.GetString($receiveBytes);
                
                # Check if the response message contains any non-printable characters
                if ( $returnData -cmatch '[^\x20-\x7F]' ) 
                {
                    # Found non-printable chars, use base64 to encode.
                    $returnData = [Convert]::ToBase64String($receiveBytes)
                }
    
                Write-Debug "This is the message you received: `"$returnData`""
                Write-Debug "This message was sent from $($RemoteIpEndPoint.Address) on their port number $($RemoteIpEndPoint.Port)."
            }
            catch [System.Exception]
            {
                Write-Debug "Caught exception -> $_"
            }
        }
    } | ForEach-Object {
        Write-Output "$(Get-Date): $($_.TotalMilliseconds)ms"
    } 
}

function Test-Server {
    param (
        [System.Net.Sockets.UdpClient] $client
    )

    $RemoteIpEndPoint = [System.Net.IPEndPoint]::new([System.Net.IPAddress]::Any, 0)
    while ($true) {
        $receiveBytes = $null
        try 
        {
            $receiveBytes = $client.Receive([ref]$RemoteIpEndPoint)
            
    
            if ($null -ne $aesKey) 
            {
                $iv = New-Object byte[] 16
                $ciphertext = New-Object byte[] ($receiveBytes.Length - 16)
                [Array]::Copy($receiveBytes, $iv, 16)
                [Array]::Copy($receiveBytes, 16, $ciphertext, 0, ($receiveBytes.Length - 16))
                $receiveBytes = $AES.DecryptCBC($ciphertext, $iv)
            } 
        }
        catch
        {
            Write-Debug "Caught Receive Exception"
        }
    
        if ($receiveBytes -gt 0)
        {
            $client.Send($receiveBytes, $receiveBytes.Length, $RemoteIpEndPoint.Address, $RemoteIpEndPoint.Port) | Write-Debug
    
            $returnData = $encoder.GetString($receiveBytes)
            if ( $returnData -cmatch '[^\x20-\x7F]' ) {
                # Found non-printable chars, use base64 to encode.
                $returnData = [Convert]::ToBase64String($receiveBytes)
            }
            
            # Uses the IPEndPoint object to determine which of these two hosts responded.
            Write-Output "{`"message`":`"$returnData`",`"address`":`"$($RemoteIpEndPoint.Address):$($RemoteIpEndPoint.Port)`"}"
        }
    }
}

try {
    $client = [System.Net.Sockets.UdpClient]::new($ListenPort)
    $client.Client.SendTimeout = $SendTimeout;
    $client.Client.ReceiveTimeout = $ReceiveTimeout;

    if ( $RunServer ) 
    {
        Test-Server $client
    }
    else
    {
        Test-Latency $client
    }
}
finally {
    $client.Close()
}

pings.sql

create table public.pings
(
    ip            inet                     not null,
    port          integer                  not null,
    message       varchar,
    created_at    timestamp with time zone default CURRENT_TIMESTAMP,
    date_reported timestamp with time zone not null,
    primary key (date_reported, ip, port)
);

alter table public.pings
    owner to postgres;

-- insert only permissions to my user
grant insert on public.pings to jsteel;