Connect-Rcon now optionally accepts a timeouts hashtable

Base class now implements the IDisposable interface

improved error handling

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")
+        $resp = $this.Send('login')
-        if ($resp -in @("Bad rcon", "Bad rconpassword.", "Invalid password.")) {
+        if ($resp -in @('Bad rcon', 'Bad rconpassword.', 'Invalid password.')) {
-            throw "invalid rcon password"
+            throw 'invalid rcon password'
         }
         $this.base.ToString() | Write-Debug
         return $this
     }
 
-    [string] Send([string]$msg) {
+    [string] Send([string]$cmd) {
-        return $this.base._send($msg)
+        $key = $cmd.Split()[0]
-    }
+        if ($this.timeouts.ContainsKey($key)) {
-
+            return $this.base._send($cmd, $this.timeouts[$key])
-    [string] Send([string]$msg, [int]$timeout) {
+        }
-        return $this.base._send($msg, $timeout)
+        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 {
+function Connect-Rcon {
-    param([string]$hostname, [int]$port, [string]$passwd)
+    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()
+    if ($rcon -and $rcon.base) {
-    "Disconnected from {0}:{1}" -f $rcon.base.hostname, $rcon.base.port | Write-Debug
+        $rcon.base.Dispose()
+        'Disconnected from {0}:{1}' -f $rcon.base.hostname, $rcon.base.port | Write-Debug
+    }
 }
 
 Export-ModuleMember -Function Connect-Rcon, Disconnect-Rcon

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)
+        $this._InitializeConnection()
-        
+    }
-        $endpoint = New-Object System.Net.IPEndPoint $ip, $this.port 
+
+    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))
+        return $this._send($msg, [Base]::DEFAULT_RECEIVE_TIMEOUT)
-
-        $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)))
     }
 
     [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 = @()
+        $this._ThrowIfDisposed()
-        $sw = [Diagnostics.Stopwatch]::StartNew()
+        
-        While ($sw.ElapsedMilliseconds -lt $timeout) {
+        if (-not $this._IsConnected()) {
-            try {
+            throw [System.InvalidOperationException]::new('Socket is not connected')
-                $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)))
+        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' ) {
+            $responseData = [System.Text.StringBuilder]::new()
-                    "finished waiting for fragment" | Write-Debug
+            $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 {
+function New-Base {
-    param([string]$hostname, [int]$port, [string]$passwd)
+    param(
+        [Parameter(Mandatory)]
+        [string]$hostname,
+        
+        [Parameter(Mandatory)]
+        [ValidateRange(1, 65535)]
+        [int]$port,
+        
+        [Parameter(Mandatory)]
+        [string]$passwd
+    )
 
     [Base]::new($hostname, $port, $passwd)
-}
+}

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()
 }