From 2800cbbcd2e150ff13a5ccfddceb9984a0e8ba4b Mon Sep 17 00:00:00 2001 From: onyx-and-iris Date: Tue, 17 Feb 2026 21:38:19 +0000 Subject: [PATCH] Connect-Rcon now optionally accepts a timeouts hashtable Base class now implements the IDisposable interface improved error handling --- lib/Q3Rcon.psm1 | 65 +++++++++------- lib/base.ps1 | 195 +++++++++++++++++++++++++++++++++++++++--------- lib/packet.ps1 | 12 +-- 3 files changed, 203 insertions(+), 69 deletions(-) diff --git a/lib/Q3Rcon.psm1 b/lib/Q3Rcon.psm1 index eccae7d..7e87be3 100644 --- a/lib/Q3Rcon.psm1 +++ b/lib/Q3Rcon.psm1 @@ -3,32 +3,41 @@ try { . (Join-Path $PSScriptRoot base.ps1) } catch { - throw "unable to dot source module files" + throw 'unable to dot source module files' } class Rcon { - [Object]$base + static [hashtable]$DefaultTimeouts = @{ + 'map' = 2000 + 'map_rotate' = 2000 + 'map_restart' = 2000 + 'fast_restart' = 2000 + } - Rcon ([string]$hostname, [int]$port, [string]$passwd) { + [Object]$base + [hashtable]$timeouts + + Rcon ([string]$hostname, [int]$port, [string]$passwd, [hashtable]$timeouts = $null) { $this.base = New-Base -hostname $hostname -port $port -passwd $passwd + $this.timeouts = $timeouts ?? [Rcon]::DefaultTimeouts } [Rcon] _login() { - $resp = $this.Send("login") - if ($resp -in @("Bad rcon", "Bad rconpassword.", "Invalid password.")) { - throw "invalid rcon password" + $resp = $this.Send('login') + if ($resp -in @('Bad rcon', 'Bad rconpassword.', 'Invalid password.')) { + throw 'invalid rcon password' } $this.base.ToString() | Write-Debug return $this } - [string] Send([string]$msg) { - return $this.base._send($msg) - } - - [string] Send([string]$msg, [int]$timeout) { - return $this.base._send($msg, $timeout) + [string] Send([string]$cmd) { + $key = $cmd.Split()[0] + if ($this.timeouts.ContainsKey($key)) { + return $this.base._send($cmd, $this.timeouts[$key]) + } + return $this.base._send($cmd) } [void] Say($msg) { @@ -36,53 +45,55 @@ class Rcon { } [void] FastRestart() { - $this.Send("fast_restart", 2000) + $this.Send('fast_restart') } [void] MapRotate() { - $this.Send("map_rotate", 2000) + $this.Send('map_rotate') } [void] MapRestart() { - $this.Send("map_restart", 2000) + $this.Send('map_restart') } [string] Map() { - return $this.Send("mapname") + return $this.Send('mapname') } [void] SetMap($mapname) { - $this.Send("map mp_" + $mapname.TrimStart("mp_"), 2000) + $this.Send('map mp_' + $mapname.TrimStart('mp_')) } [string] Gametype() { - return $this.Send("g_gametype") + return $this.Send('g_gametype') } [void] SetGametype($gametype) { - $this.Send("g_gametype $gametype") + $this.Send('g_gametype', $gametype) } [string] HostName() { - return $this.Send("sv_hostname") + return $this.Send('sv_hostname') } [void] SetHostName($hostname) { - $this.Send("sv_hostname $hostname") + $this.Send('sv_hostname', $hostname) } } -Function Connect-Rcon { - param([string]$hostname, [int]$port, [string]$passwd) +function Connect-Rcon { + param([string]$hostname, [int]$port, [string]$passwd, [Parameter(Mandatory = $false)][hashtable]$timeouts) - [Rcon]::new($hostname, $port, $passwd)._login() + [Rcon]::new($hostname, $port, $passwd, $timeouts)._login() } -Function Disconnect-Rcon { +function Disconnect-Rcon { param([Rcon]$rcon) - $rcon.base._close() - "Disconnected from {0}:{1}" -f $rcon.base.hostname, $rcon.base.port | Write-Debug + if ($rcon -and $rcon.base) { + $rcon.base.Dispose() + 'Disconnected from {0}:{1}' -f $rcon.base.hostname, $rcon.base.port | Write-Debug + } } Export-ModuleMember -Function Connect-Rcon, Disconnect-Rcon diff --git a/lib/base.ps1 b/lib/base.ps1 index 7b065ad..87d4a6c 100644 --- a/lib/base.ps1 +++ b/lib/base.ps1 @@ -1,83 +1,206 @@ class Base { + static [int]$DEFAULT_RECEIVE_TIMEOUT = 100 + static [int]$DEFAULT_BUFFER_SIZE = 4096 + static [int]$DEFAULT_SEND_TIMEOUT = 5000 + [string]$hostname [int]$port [string]$passwd [Object]$request [Object]$response hidden [System.Net.Sockets.Socket] $_socket + hidden [bool]$_disposed = $false + hidden [byte[]]$_receiveBuffer Base ([string]$hostname, [int]$port, [string]$passwd) { + if ([string]::IsNullOrWhiteSpace($hostname)) { + throw [System.ArgumentException]::new('Hostname cannot be null or empty', 'hostname') + } + if ($port -le 0 -or $port -gt 65535) { + throw [System.ArgumentOutOfRangeException]::new('port', 'Port must be between 1-65535') + } + if ([string]::IsNullOrWhiteSpace($passwd)) { + throw [System.ArgumentException]::new('Password cannot be null or empty', 'passwd') + } + $this.hostname = $hostname $this.port = $port $this.passwd = $passwd $this.request = New-RequestPacket($this.passwd) $this.response = New-ResponsePacket + $this._receiveBuffer = [byte[]]::new([Base]::DEFAULT_BUFFER_SIZE) - $ip = [system.net.IPAddress]::Parse([System.Net.Dns]::GetHostAddresses($this.hostname)[0].IPAddressToString) - - $endpoint = New-Object System.Net.IPEndPoint $ip, $this.port + $this._InitializeConnection() + } + + hidden [void] _InitializeConnection() { try { - $this._socket = [System.Net.Sockets.Socket]::New( + $hostEntry = [System.Net.Dns]::GetHostEntry($this.hostname) + if ($hostEntry.AddressList.Length -eq 0) { + throw [System.Net.Sockets.SocketException]::new([int][System.Net.Sockets.SocketError]::HostNotFound) + } + + $ipv4Address = $hostEntry.AddressList | Where-Object { $_.AddressFamily -eq [System.Net.Sockets.AddressFamily]::InterNetwork } | Select-Object -First 1 + if (-not $ipv4Address) { + throw [System.InvalidOperationException]::new("No IPv4 address found for hostname: $($this.hostname)") + } + + $endpoint = [System.Net.IPEndPoint]::new($ipv4Address, $this.port) + + $this._socket = [System.Net.Sockets.Socket]::new( [System.Net.Sockets.AddressFamily]::InterNetwork, [System.Net.Sockets.SocketType]::Dgram, [System.Net.Sockets.ProtocolType]::UDP ) + $this._socket.Connect($endpoint) - $this._socket.ReceiveTimeout = 100 + $this._socket.ReceiveTimeout = [Base]::DEFAULT_RECEIVE_TIMEOUT + $this._socket.SendTimeout = [Base]::DEFAULT_SEND_TIMEOUT } catch [System.Net.Sockets.SocketException] { - throw "Failed to create UDP connection to server." + $this._Cleanup() + throw [System.InvalidOperationException]::new( + "Failed to create UDP connection to $($this.hostname):$($this.port). Error: $($_.Exception.Message)", + $_.Exception + ) } + catch { + $this._Cleanup() + throw [System.InvalidOperationException]::new( + "Failed to initialize connection to $($this.hostname):$($this.port). Error: $($_.Exception.Message)", + $_.Exception + ) + } + } + + hidden [void] _ThrowIfDisposed() { + if ($this._disposed) { + throw [System.ObjectDisposedException]::new($this.GetType().Name) + } + } + + hidden [bool] _IsConnected() { + return $this._socket -and -not $this._disposed -and $this._socket.Connected } [string] ToString () { - return "Rcon connection {0}:{1} with pass {2}" -f $this.hostname, $this.port, $this.passwd + $status = if ($this._IsConnected()) { 'Connected' } else { 'Disconnected' } + return 'Rcon connection {0}:{1} ({2})' -f $this.hostname, $this.port, $status } [string] _send([string]$msg) { - $this._socket.Send($this.request.Payload($msg)) - - $buf = New-Object System.Byte[] 4096 - try { - $this._socket.Receive($buf) - } - catch [System.Net.Sockets.SocketException] { - if ( $_.Exception.SocketErrorCode -eq 'TimedOut' ) { - "finished waiting for fragment" | Write-Debug - } - } - return [System.Text.Encoding]::ASCII.GetString($($buf | Select-Object -Skip $($this.response.Header().Length - 1))) + return $this._send($msg, [Base]::DEFAULT_RECEIVE_TIMEOUT) } [string] _send([string]$msg, [int]$timeout) { - $this._socket.Send($this.request.Payload($msg)) + if ([string]::IsNullOrEmpty($msg)) { + throw [System.ArgumentException]::new('Message cannot be null or empty', 'msg') + } + if ($timeout -le 0) { + throw [System.ArgumentOutOfRangeException]::new('timeout', 'Timeout must be positive') + } - [string[]]$data = @() - $sw = [Diagnostics.Stopwatch]::StartNew() - While ($sw.ElapsedMilliseconds -lt $timeout) { - try { - $buf = New-Object System.Byte[] 4096 - $this._socket.Receive($buf) - $data += [System.Text.Encoding]::ASCII.GetString($($buf | Select-Object -Skip $($this.response.Header().Length - 1))) + $this._ThrowIfDisposed() + + if (-not $this._IsConnected()) { + throw [System.InvalidOperationException]::new('Socket is not connected') + } + + try { + $payload = $this.request.Payload($msg) + $bytesSent = $this._socket.Send($payload) + if ($bytesSent -ne $payload.Length) { + Write-Warning "Not all bytes were sent. Expected: $($payload.Length), Sent: $bytesSent" } - catch [System.Net.Sockets.SocketException] { - if ( $_.Exception.SocketErrorCode -eq 'TimedOut' ) { - "finished waiting for fragment" | Write-Debug + + $responseData = [System.Text.StringBuilder]::new() + $sw = [System.Diagnostics.Stopwatch]::StartNew() + $headerLength = $this.response.Header().Length + + do { + try { + $bytesReceived = $this._socket.Receive($this._receiveBuffer) + if ($bytesReceived -gt 0) { + $dataStartIndex = [Math]::Min($headerLength, $bytesReceived) + $responseText = [System.Text.Encoding]::ASCII.GetString($this._receiveBuffer, $dataStartIndex, $bytesReceived - $dataStartIndex) + $responseData.Append($responseText) | Out-Null + } + } + catch [System.Net.Sockets.SocketException] { + if ($_.Exception.SocketErrorCode -eq 'TimedOut') { + Write-Debug 'Socket receive timeout - continuing to wait for more data' + continue + } + else { + throw [System.InvalidOperationException]::new( + "Socket error during receive: $($_.Exception.Message)", + $_.Exception + ) + } + } + } while ($sw.ElapsedMilliseconds -lt $timeout) + + $sw.Stop() + return $responseData.ToString() + } + catch [System.Net.Sockets.SocketException] { + throw [System.InvalidOperationException]::new( + "Network error during send/receive: $($_.Exception.Message)", + $_.Exception + ) + } + } + + hidden [void] _Cleanup() { + if ($this._socket) { + try { + if ($this._socket.Connected) { + $this._socket.Shutdown([System.Net.Sockets.SocketShutdown]::Both) } } + catch { + Write-Debug "Error during socket shutdown: $($_.Exception.Message)" + } + finally { + $this._socket.Close() + $this._socket = $null + } } - $sw.Stop() - return [string]::Join("", $data) } [void] _close() { - $this._socket.Close() + $this.Dispose() + } + + # Dispose implementation (following IDisposable pattern) + [void] Dispose() { + $this.Dispose($true) + [System.GC]::SuppressFinalize($this) + } + + hidden [void] Dispose([bool]$disposing) { + if (-not $this._disposed) { + if ($disposing) { + $this._Cleanup() + } + $this._disposed = $true + } } } -Function New-Base { - param([string]$hostname, [int]$port, [string]$passwd) +function New-Base { + param( + [Parameter(Mandatory)] + [string]$hostname, + + [Parameter(Mandatory)] + [ValidateRange(1, 65535)] + [int]$port, + + [Parameter(Mandatory)] + [string]$passwd + ) [Base]::new($hostname, $port, $passwd) -} +} \ No newline at end of file diff --git a/lib/packet.ps1 b/lib/packet.ps1 index acf1b1c..c384fc2 100644 --- a/lib/packet.ps1 +++ b/lib/packet.ps1 @@ -2,7 +2,7 @@ class Packet { [System.Byte[]]$MAGIC = @(, 0xFF * 4) [string] Header() { - throw "method not implemented" + throw 'method not implemented' } } @@ -14,24 +14,24 @@ class RequestPacket : Packet { } [System.Byte[]] Header() { - return $this.MAGIC + [System.Text.Encoding]::ASCII.GetBytes("rcon") + return $this.MAGIC + [System.Text.Encoding]::ASCII.GetBytes('rcon') } [System.Byte[]] Payload([string]$msg) { - return $this.Header() + [System.Text.Encoding]::ASCII.GetBytes($(" {0} {1}" -f $this.passwd, $msg)) + return $this.Header() + [System.Text.Encoding]::ASCII.GetBytes($(' {0} {1}' -f $this.passwd, $msg)) } } class ResponsePacket : Packet { [System.Byte[]] Header() { - return $this.MAGIC + [System.Text.Encoding]::ASCII.GetBytes("print\n") + return $this.MAGIC + [System.Text.Encoding]::ASCII.GetBytes('print\n') } } -Function New-RequestPacket([string]$passwd) { +function New-RequestPacket([string]$passwd) { [RequestPacket]::new($passwd) } -Function New-ResponsePacket { +function New-ResponsePacket { [ResponsePacket]::new() }