lib/ionian/socket.rb in ionian-0.4.0 vs lib/ionian/socket.rb in ionian-0.4.1

@@ -1,27 +1,28 @@
 require 'ionian/extension/socket'
 
 module Ionian
+  
+  # A convenient wrapper for TCP, UDP, and Unix sockets.
   class Socket
     
-    ############
-    # TODO NOTES
-    ############
-    # Always lazily instiantiate @socket, even when persistent?
-    # May not work with forwarding method calls.
-    # Oh! Unless the forwarded methods check for @socket to exist.
-    # Will persistent methods have to check for the socket not to be
-    # closed as well?
-    
-    def initialize(**kvargs)
+    # Args:
+    #   host: IP or hostname to connect to.
+    #   port: Connection's port number. Default is 23. Unused by :unix protocol.
+    #   protocol: Type of socket to create. :tcp, :udp, :unix. Default is :tcp.
+    #   persistent: The socket remains open after data is sent if this is true.
+    #               The socket closes after data is sent and a packet is received
+    #               if this is false. Default is true.
+    #   expression: Overrides the #read_match regular expression for received data.
+    def initialize(**kwargs)
       @socket         = nil
       
-      @host           = kvargs.fetch :host
-      @port           = kvargs.fetch :port, 23
-      @expression     = kvargs.fetch :expression, nil
-      @protocol       = kvargs.fetch :protocol, :tcp
-      @persistent     = kvargs.fetch :persistent, true
+      @host           = kwargs.fetch :host
+      @port           = kwargs.fetch :port, 23
+      @protocol       = kwargs.fetch :protocol, :tcp
+      @persistent     = kwargs.fetch :persistent, true
+      @expression     = kwargs.fetch :expression, nil
       
       create_socket if @persistent
     end
         
     # Returns a symbol of the type of protocol this socket uses:
@@ -33,19 +34,26 @@
     # Returns true if the socket remains open after writing data.
     def persistent?
       @persistent == false || @persistent == nil ? false : true
     end
     
-    # Send a command (data) to the socket. Returns received matches.
+    # Send a command (data) to the socket.
+    # Returns an array of received matches.
     # Block yields received match.
-    # See Ionian::Extension::IO#read_match
-    def cmd(string, **kvargs, &block)
+    # See Ionian::Extension::IO#read_match.
+    def cmd(string, **kwargs, &block)
       create_socket unless @persistent
-      @socket.write string
+      
+      if @protocol == :udp
+        @socket.send string, 0
+      else
+        @socket.write string
+      end
+      
       @socket.flush
       
-      matches = @socket.read_match(kvargs) {|match| yield match if block_given?}
+      matches = @socket.read_match(kwargs) {|match| yield match if block_given?}
       @socket.close unless @persistent
       
       matches
     end
     
@@ -54,13 +62,13 @@
     
     # Returns true if there is data in the receive buffer.
     # Args:
     #   Timeout: Number of seconds to wait for data until
     #     giving up. Set to nil for blocking.
-    def has_data?(**kvargs)
+    def has_data?(**kwargs)
       return false unless @socket
-      @socket.has_data? kvargs
+      @socket.has_data? kwargs
     end
     
     # Returns true if the socket is closed.
     def closed?
       return true unless @socket
@@ -81,12 +89,19 @@
     
     # Writes the given string to the socket. Returns the number of
     # bytes written.
     def write(string)
       create_socket unless @persistent
-      num_bytes = @socket.write string
       
+      num_bytes = 0
+      
+      if @protocol == :udp
+        num_bytes = @socket.send string, 0
+      else
+        num_bytes = @socket.write string
+      end
+      
       unless @persistent
         # Read in data to prevent RST packets.
         has_data = ::IO.select [@socket], nil, nil, 0
         @socket.readpartial 0xFFFF if has_data
         
@@ -99,17 +114,19 @@
     alias_method :<<, :write
     
     
     private
     
+    # Initialize or reinitialize @socket.
     def create_socket
       @socket.close if @socket and not @socket.closed?
       
       case @protocol
       when :tcp
         @socket = ::TCPSocket.new @host, @port
       when :udp
         @socket = ::UDPSocket.new
+        @socket.bind '', @port
         @socket.connect @host, @port
       when :unix
         @socket = ::UNIXSocket.new @host
       end
       
\ No newline at end of file