basic tcp test

git-svn-id: file:///home/svn/incoming/trunk@2547 4d416f70-5f16-0410-b530-b9f4589650da

lib/rex.rb.ts.rb

@@ -3,6 +3,7 @@
 require 'test/unit'
 require 'Rex'
 require 'Rex/Socket.rb.ut'
+require 'Rex/Socket/Tcp.rb.ut'
 require 'Rex/Socket/Parameters.rb.ut'
 require 'Rex/Socket/Comm/Local.rb.ut'
 
@@ -12,6 +13,7 @@ class Rex::TestSuite
 
 		suite << Rex::Socket::UnitTest.suite
 		suite << Rex::Socket::Parameters::UnitTest.suite
+		suite << Rex::Socket::Tcp::UnitTest.suite
 		suite << Rex::Socket::Comm::Local::UnitTest.suite
 
 		return suite;

lib/rex/io/stream.rb

@@ -88,8 +88,6 @@ module Stream
 			timeout(wait) {
 				return write(buf, opts)
 			}
-
-			raise TimeoutError, "Write operation timed out.", caller
 		else
 			return write(buf, opts)
 		end
@@ -103,8 +101,6 @@ module Stream
 			timeout(wait) {
 				return read(length, opts)
 			}
-			
-			raise TimeoutError, "Read operation timed out.", caller
 		else
 			return read(length, opts)
 		end
@@ -142,7 +138,7 @@ module Stream
 
 		# Keep looping until there is no more data to be gotten..
 		while (poll_read(ltimeout))
-			temp = recv(def_block_size)
+			temp = read(def_block_size)
 
 			break if (temp.empty?)
 

lib/rex/socket/tcp.rb

@@ -42,8 +42,8 @@ class Rex::Socket::Tcp < Rex::Socket
 		end
 	end
 
-	def shutdown(how = SW_BOTH)
+	def shutdown(how = SHUT_RDWR)
-		return sock.shutdown(how)
+		return (sock.shutdown(how) == 0)
 	end
 
 	def close

lib/rex/socket/tcp.rb.ut.rb

@@ -0,0 +1,47 @@
+#!/usr/bin/ruby
+
+$:.unshift(File.join('..', '..', File.dirname(__FILE__)))
+
+require 'test/unit'
+require 'Rex/Socket/Tcp'
+
+class Rex::Socket::Tcp::UnitTest < Test::Unit::TestCase
+
+	def test_tcp
+		serv_port = 65432
+		serv = TCPServer.new('127.0.0.1', serv_port)
+
+		begin
+			# Connect to the temp server
+			t = nil
+			assert_nothing_raised {
+				t = Rex::Socket.create_tcp(
+					'PeerHost' => '127.0.0.1',
+					'PeerPort' => serv_port)
+			}
+			assert_kind_of(Rex::Socket::Tcp, t, "valid tcp socket")
+
+			# Accept the client connection
+			serv_con = serv.accept
+			assert_kind_of(TCPSocket, serv_con, "valid server socket connection")
+
+			assert_equal(5, t.write("test\n"), "cli: write test")
+			assert_equal("test\n", serv_con.recv(5), "srv: read test")
+			assert_equal(10, serv_con.send("A" * 10, 10), "srv: write A*10")
+			assert_equal("A" * 10, t.get, "cli: gobble A*10")
+			assert_equal(5, t << "test\n", "cli: << test")
+			assert_equal("test\n", serv_con.recv(5), "srv: read test (2)")
+			assert_equal(5, serv_con.send("testa", 6), "srv: write testa (3)")
+			assert_kind_of(Array, t.poll_read(1), "cli: poll read")
+			assert_equal("testa", t.get, "cli: gobble testa")
+			assert_equal(true, t.shutdown(::Socket::SHUT_RD), "cli: shutdown read")
+			assert_equal(true, t.shutdown(::Socket::SHUT_WR), "cli: shutdown read")
+			assert_nothing_raised {
+				t.close
+			}
+		ensure
+			serv.close
+		end
+	end
+
+end