LatencyDiag.psm1

Set-StrictMode -Version Latest
$ErrorActionPreference = 'Stop'

function _NowIso { (Get-Date).ToString('s') }

function _Fail {
  param(
    [Parameter(Mandatory)][string]$Message,
    [string]$Hint
  )
  $full = if ($Hint) { "$Message`nHINT: $Hint" } else { $Message }
  throw ([System.ArgumentException]::new($full))
}

function _EnsureDirectory {
  param([string]$Path, [string]$Purpose = "output")
  if ([string]::IsNullOrWhiteSpace($Path)) { return }
  if (-not (Test-Path -LiteralPath $Path)) {
    try { New-Item -ItemType Directory -Path $Path -Force | Out-Null }
    catch { _Fail "Cannot create $Purpose folder: $Path" "Check permissions and that the path is valid." }
  }
}

function _Stopwatch([ScriptBlock]$Script) {
  $sw = [System.Diagnostics.Stopwatch]::StartNew()
  $ok = $true
  try { & $Script | Out-Null } catch { $ok = $false } finally { $sw.Stop() }
  [pscustomobject]@{ Success = $ok; ElapsedMs = [double]$sw.ElapsedMilliseconds }
}

function _ComputeStats {
  param(
    [double[]] $SamplesMs = @(),
    [int] $Iterations = 0,
    [int] $FailCount = 0
  )

  if ($null -eq $SamplesMs) { $SamplesMs = @() }

  $ok = @($SamplesMs | Where-Object { $_ -ge 0 })
  $successCount = $ok.Count
  $lossPct = if ($Iterations -gt 0) { [math]::Round(($FailCount / $Iterations) * 100, 2) } else { 0 }

  if ($successCount -eq 0) {
    return [pscustomobject]@{
      Iterations    = $Iterations
      SuccessCount  = 0
      FailCount     = $FailCount
      LossPct       = $lossPct
      MinMs         = $null
      AvgMs         = $null
      P50Ms         = $null
      P95Ms         = $null
      P99Ms         = $null
      MaxMs         = $null
      JitterMs      = $null
    }
  }

  $sorted = $ok | Sort-Object
  function _Percentile([double[]]$arr, [double]$p) {
    if ($arr.Count -eq 0) { return $null }
    $idx = [math]::Ceiling(($p/100) * $arr.Count) - 1
    if ($idx -lt 0) { $idx = 0 }
    if ($idx -ge $arr.Count) { $idx = $arr.Count - 1 }
    [double]$arr[$idx]
  }

  $avg = ($ok | Measure-Object -Average).Average
  $min = $sorted[0]
  $max = $sorted[-1]
  $p50 = _Percentile $sorted 50
  $p95 = _Percentile $sorted 95
  $p99 = _Percentile $sorted 99

  $mean = [double]$avg
  $variance = 0.0
  foreach ($v in $ok) { $variance += [math]::Pow(($v - $mean), 2) }
  $variance = $variance / [math]::Max(1, $ok.Count)
  $jitter = [math]::Sqrt($variance)

  [pscustomobject]@{
    Iterations    = $Iterations
    SuccessCount  = $successCount
    FailCount     = $FailCount
    LossPct       = $lossPct
    MinMs         = [math]::Round($min, 2)
    AvgMs         = [math]::Round($avg, 2)
    P50Ms         = [math]::Round($p50, 2)
    P95Ms         = [math]::Round($p95, 2)
    P99Ms         = [math]::Round($p99, 2)
    MaxMs         = [math]::Round($max, 2)
    JitterMs      = [math]::Round($jitter, 2)
  }
}

function _RedactValue {
  param([string]$Value, [switch]$Redact)
  if (-not $Redact) { return $Value }
  if ([string]::IsNullOrWhiteSpace($Value)) { return $Value }
  $keep = 6
  if ($Value.Length -le $keep) { return ('*' * $Value.Length) }
  ('*' * ($Value.Length - $keep)) + $Value.Substring($Value.Length - $keep)
}

function _TcpConnectOnce {
  param([Parameter(Mandatory)][string]$Host,[Parameter(Mandatory)][int]$Port,[int]$TimeoutMs = 2000)
  $sw = [System.Diagnostics.Stopwatch]::StartNew()
  $client = New-Object System.Net.Sockets.TcpClient
  try {
    $iar = $client.BeginConnect($Host, $Port, $null, $null)
    if (-not $iar.AsyncWaitHandle.WaitOne($TimeoutMs, $false)) {
      $client.Close()
      throw "TCP connect timeout after ${TimeoutMs}ms"
    }
    $client.EndConnect($iar)
    $sw.Stop()
    return [double]$sw.ElapsedMilliseconds
  } finally { $client.Dispose() }
}

function _HttpOnce {
  param([Parameter(Mandatory)][string]$Url,[int]$TimeoutMs = 4000,[switch]$NoDownload)

  $handler = New-Object System.Net.Http.HttpClientHandler
  $client  = New-Object System.Net.Http.HttpClient($handler)
  $client.Timeout = [TimeSpan]::FromMilliseconds($TimeoutMs)

  $sw = [System.Diagnostics.Stopwatch]::StartNew()
  try {
    $method = if ($NoDownload) { [System.Net.Http.HttpMethod]::Head } else { [System.Net.Http.HttpMethod]::Get }
    $req = New-Object System.Net.Http.HttpRequestMessage($method, $Url)
    $resp = $client.SendAsync($req, [System.Net.Http.HttpCompletionOption]::ResponseHeadersRead).GetAwaiter().GetResult()
    $resp.EnsureSuccessStatusCode() | Out-Null

    if (-not $NoDownload -and $method -eq [System.Net.Http.HttpMethod]::Get) {
      $stream = $resp.Content.ReadAsStreamAsync().GetAwaiter().GetResult()
      $buffer = New-Object byte[] (32768)
      [void]$stream.Read($buffer, 0, $buffer.Length)
    }

    $sw.Stop()
    return [double]$sw.ElapsedMilliseconds
  } finally { $client.Dispose(); $handler.Dispose() }
}

function _BytesToMegaBitsPerSec([int64]$Bytes, [double]$Seconds) {
  if ($Seconds -le 0) { return $null }
  [math]::Round((($Bytes * 8.0) / $Seconds) / 1e6, 2)
}

function _BytesToMegaBytesPerSec([int64]$Bytes, [double]$Seconds) {
  if ($Seconds -le 0) { return $null }
  [math]::Round(($Bytes / 1e6) / $Seconds, 2)
}

function Test-NetLatency {
  [CmdletBinding()]
  param(
    [Parameter(Mandatory)] [string[]] $Targets,
    [int] $Iterations = 20,
    [int] $TimeoutMs = 2000,
    [int[]] $TcpPorts = @(443),
    [string[]] $HttpUrls = @(),
    [switch] $NoDownload,
    [switch] $Redact
  )

  if ($Iterations -le 0) { _Fail "Iterations must be > 0." "Example: -Iterations 20" }
  if ($TimeoutMs -lt 200) { _Fail "TimeoutMs is too low." "Use at least 500–2000ms." }

  $results = New-Object System.Collections.Generic.List[object]

  foreach ($t in $Targets) {
    if ([string]::IsNullOrWhiteSpace($t)) { continue }
    $targetLabel = _RedactValue $t -Redact:$Redact

    $icmpSamples = @()
    $icmpFails = 0
    for ($i=0; $i -lt $Iterations; $i++) {
      try {
        $ms = $null
        if (Get-Command Test-Connection -ErrorAction SilentlyContinue) {
          $r = Test-Connection -TargetName $t -Count 1 -TimeoutSeconds ([math]::Ceiling($TimeoutMs/1000)) -ErrorAction Stop
          $ms = @($r | Select-Object -First 1 -ExpandProperty ResponseTime -ErrorAction SilentlyContinue)
          if ($null -eq $ms) { $ms = @($r | Select-Object -First 1 -ExpandProperty Latency -ErrorAction SilentlyContinue) }
        }
        if ($null -eq $ms) { throw "No ICMP latency available" }
        $icmpSamples += [double]$ms
      } catch { $icmpFails++ }
    }

    $results.Add([pscustomobject]@{
      TestType    = 'ICMP'
      Target      = $targetLabel
      Timestamp   = _NowIso
      Stats       = _ComputeStats -SamplesMs $icmpSamples -Iterations $Iterations -FailCount $icmpFails
      Notes       = if ($icmpFails -gt 0) { 'Some ICMP failed (ICMP may be blocked). Compare with TCP.' } else { 'ICMP is RTT (round-trip).' }
    }) | Out-Null

    foreach ($p in $TcpPorts) {
      if ($p -lt 1 -or $p -gt 65535) { _Fail "Invalid TCP port: $p" "Use 1..65535." }

      $tcpSamples = @()
      $tcpFails = 0
      for ($i=0; $i -lt $Iterations; $i++) {
        try { $tcpSamples += (_TcpConnectOnce -Host $t -Port $p -TimeoutMs $TimeoutMs) } catch { $tcpFails++ }
      }

      $results.Add([pscustomobject]@{
        TestType  = 'TCP'
        Target    = "${targetLabel}:$p"
        Timestamp = _NowIso
        Stats     = _ComputeStats -SamplesMs $tcpSamples -Iterations $Iterations -FailCount $tcpFails
        Notes     = if ($tcpFails -gt 0) { "Some TCP connects failed. Check firewall/DNS. TimeoutMs=$TimeoutMs" } else { 'TCP connect = reachability + handshake.' }
      }) | Out-Null
    }
  }

  foreach ($u in $HttpUrls) {
    if ([string]::IsNullOrWhiteSpace($u)) { continue }
    $urlLabel = _RedactValue $u -Redact:$Redact
    $httpSamples = @()
    $httpFails = 0
    for ($i=0; $i -lt $Iterations; $i++) {
      try { $httpSamples += (_HttpOnce -Url $u -TimeoutMs $TimeoutMs -NoDownload:$NoDownload) } catch { $httpFails++ }
    }
    $results.Add([pscustomobject]@{
      TestType  = 'HTTP'
      Target    = $urlLabel
      Timestamp = _NowIso
      Stats     = _ComputeStats -SamplesMs $httpSamples -Iterations $Iterations -FailCount $httpFails
      Notes     = if ($NoDownload) { 'HEAD/headers-only (no body).' } else { 'GET with tiny (<=32KB) read.' }
    }) | Out-Null
  }

  return $results
}

# Minimal HTTP throughput test (real download)
function Test-HttpTransfer {
  [CmdletBinding()]
  param(
    [Parameter(Mandatory)][string]$DownloadUrl,
    [int]$TimeoutMs = 300000,
    [switch]$Redact
  )

  $urlLabel = _RedactValue $DownloadUrl -Redact:$Redact

  $handler = New-Object System.Net.Http.HttpClientHandler
  $client  = New-Object System.Net.Http.HttpClient($handler)
  $client.Timeout = [TimeSpan]::FromMilliseconds($TimeoutMs)

  $bytesRead = [int64]0
  $sw = [System.Diagnostics.Stopwatch]::StartNew()

  try {
    $req = New-Object System.Net.Http.HttpRequestMessage([System.Net.Http.HttpMethod]::Get, $DownloadUrl)
    $resp = $client.SendAsync($req, [System.Net.Http.HttpCompletionOption]::ResponseHeadersRead).GetAwaiter().GetResult()
    $resp.EnsureSuccessStatusCode() | Out-Null

    $stream = $resp.Content.ReadAsStreamAsync().GetAwaiter().GetResult()
    $buffer = New-Object byte[] (131072)
    $maxBytes = 50MB

    while ($true) {
      $read = $stream.Read($buffer, 0, $buffer.Length)
      if ($read -le 0) { break }
      $bytesRead += $read
      if ($bytesRead -ge $maxBytes) { break }
    }

    $sw.Stop()
    $secs = [math]::Max(0.001, $sw.Elapsed.TotalSeconds)

    return [pscustomobject]@{
      TestType         = 'HTTP-Download'
      Target           = $urlLabel
      Timestamp        = _NowIso
      Bytes            = $bytesRead
      Seconds          = [math]::Round($secs, 3)
      MegaBytesPerSec  = _BytesToMegaBytesPerSec -Bytes $bytesRead -Seconds $secs
      MegaBitsPerSec   = _BytesToMegaBitsPerSec  -Bytes $bytesRead -Seconds $secs
      Notes            = "Real download sample up to 50MB."
    }
  } finally {
    $client.Dispose(); $handler.Dispose()
  }
}

function Invoke-LatencyDiag {
  [CmdletBinding()]
  param(
    [string[]] $InternetTargets = @('1.1.1.1','8.8.8.8'),
    [int[]] $TcpPorts = @(443,53,445),
    [string[]] $HttpUrls = @('https://www.microsoft.com'),
    [int] $Iterations = 10,
    [int] $TimeoutMs = 2500,
    [switch] $NoDownload,
    [switch] $Redact,
    [switch] $PassThru
  )

  $all = New-Object System.Collections.Generic.List[object]
  $all.AddRange(@(Test-NetLatency -Targets $InternetTargets -Iterations $Iterations -TimeoutMs $TimeoutMs -TcpPorts $TcpPorts -HttpUrls $HttpUrls -NoDownload:$NoDownload -Redact:$Redact)) | Out-Null

  if ($PassThru) { return $all.ToArray() }

  $all.ToArray() |
    Where-Object { $_.Stats } |
    Select-Object TestType, Target,
      @{n='Loss%';e={$_.Stats.LossPct}},
      @{n='P50(ms)';e={$_.Stats.P50Ms}},
      @{n='P95(ms)';e={$_.Stats.P95Ms}},
      @{n='P99(ms)';e={$_.Stats.P99Ms}},
      @{n='Avg(ms)';e={$_.Stats.AvgMs}} |
    Sort-Object TestType, Target |
    Format-Table -AutoSize
}

Export-ModuleMember -Function Invoke-LatencyDiag, Test-NetLatency, Test-HttpTransfer