add Socksify.proxy method

code-later · Mar 31, 2011 · ae84017 · ae84017
1 parent 5b5bfe1
commit ae84017
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 0 deletions.
diff --git a/lib/socksify.rb b/lib/socksify.rb
@@ -298,4 +298,17 @@ def self.resolve(host)
       s.close
     end
   end
+
+  def self.proxy(server, port)
+    default_server = TCPSocket::socks_server
+    default_port = TCPSocket::socks_port
+    begin
+      TCPSocket::socks_server = server
+      TCPSocket::socks_port = port
+      yield
+    ensure
+      TCPSocket::socks_server = default_server
+      TCPSocket::socks_port = default_port
+    end
+  end
 end
diff --git a/test/tc_socksify.rb b/test/tc_socksify.rb
@@ -119,4 +119,38 @@ def test_resolve_reverse
       Socksify::resolve("0.0.0.0")
     end
   end
+
+  def test_proxy
+    enable_socks 
+
+    default_server = TCPSocket.socks_server
+    default_port = TCPSocket.socks_port
+
+    Socksify.proxy('localhost.example.com', 60001) {
+      assert_equal TCPSocket.socks_server, 'localhost.example.com'
+      assert_equal TCPSocket.socks_port, 60001
+    }
+
+    assert_equal TCPSocket.socks_server, default_server
+    assert_equal TCPSocket.socks_port, default_port
+  end
+
+  def test_proxy_failback
+    enable_socks 
+
+    default_server = TCPSocket.socks_server
+    default_port = TCPSocket.socks_port
+
+    assert_raise StandardError do
+      Socksify.proxy('localhost.example.com', 60001) {
+        raise StandardError.new('error')
+      }
+    end
+
+    assert_equal TCPSocket.socks_server, default_server
+    assert_equal TCPSocket.socks_port, default_port
+  end
 end
+
+
+