add udp and client src and specs

jkthorne · Nov 7, 2018 · 0564bc4 · 0564bc4
1 parent e4cd539
commit 0564bc4
Show file tree

Hide file tree

Showing 7 changed files with 115 additions and 14 deletions.
diff --git a/README.md b/README.md
@@ -58,6 +58,9 @@ ssh -D 1080 -C -N [email protected]
 ```
 
 ### Basic Usage
+
+#### Basic Socket
+
 To open a socks connection and send a basic HTTP request and get a response.
 
 ```crystal
@@ -75,6 +78,49 @@ if response.success?
 end
 ```
 
+#### Basic Client
+
+`Socks::Client` functions almost like the Crystal (HTTP::Client)[https://crystal-lang.org/api/latest/HTTP/Client.html]
+
+```crystal
+client = Socks::Client.new("www.example.com", host_addr: "127.0.0.1", host_port: 1080)
+response = client.get("/")
+puts response.status_code      # => 200
+puts response.body.lines.first # => "<!doctype html>"
+client.close
+```
+
+#### Basic UDP
+
+`Socks::UDP` functions almost like the Crystal (UDPSocket)[https://crystal-lang.org/api/latest/UDPSocket.html]
+
+```crystal
+server = UDPSocket.new
+server.bind "localhost", 9999
+
+client = Socks::UDP.new(host_addr: "127.0.0.1", host_port: 1080)
+client.connect "localhost", 9999
+
+client.send "message"
+message, client_addr = server.receive
+
+message     # => "message"
+client_addr # => Socket::IPAddress(127.0.0.1:50516)
+
+client.close
+server.close
+```
+
+you can use `Socks::Client` almost like the Crystal (HTTP::Client)[https://crystal-lang.org/api/latest/HTTP/Client.html]
+
+```crystal
+client = Socks::Client.new("www.example.com", host_addr: "127.0.0.1", host_port: 1080)
+response = client.get("/")
+puts response.status_code      # => 200
+puts response.body.lines.first # => "<!doctype html>"
+client.close
+```
+
 ### Remote server connnection
 
 To open a connection to a remote SOCKS5 Server.

diff --git a/spec/socks/client_spec.cr b/spec/socks/client_spec.cr
@@ -0,0 +1,22 @@
+require "../spec_helper"
+
+describe Socks::Client do
+  it "connect" do
+    begin
+      server = HTTP::Server.new do |context|
+        context.response.content_type = "text/plain"
+        if context.request.path == "/ping"
+          context.response << "pong"
+        end
+      end
+      address = server.bind_unused_port "127.0.0.1"
+      spawn { server.try &.listen }
+
+      client = Socks::Client.new("127.0.0.1", address.port, host_addr: "127.0.0.1", host_port: SSH_PORT)
+      response = client.get("/ping")
+      response.try &.body.should eq "pong"
+    ensure
+      server.try &.close
+    end
+  end
+end
diff --git a/spec/socks/udp_spec.cr b/spec/socks/udp_spec.cr
@@ -0,0 +1,23 @@
+require "../spec_helper"
+
+describe Socks::UDP do
+  it "associate" do
+    begin
+      udp_port = rand(8000..10000)
+      server = UDPSocket.new
+      server.bind "localhost", udp_port
+
+
+      client = Socks::UDP.new(host_addr: "127.0.0.1", host_port: SSH_PORT)
+      client.connect("localhost", udp_port)
+
+      client.send("yolo")
+      message, client_addr = server.receive
+
+      message.should eq "yolo"
+    ensure
+      client.try &.close
+      server.try &.close
+    end
+  end
+end
diff --git a/spec/socks_spec.cr b/spec/socks_spec.cr
@@ -21,7 +21,7 @@ describe Socks do
       server = HTTP::Server.new do |context|
         context.response.content_type = "text/plain"
         if context.request.path == "/ping"
-          context.response.puts "pong"
+          context.response << "pong"
         end
       end
       address = server.bind_unused_port "127.0.0.1"
@@ -37,7 +37,7 @@ describe Socks do
       socket.flush
       response = HTTP::Client::Response.from_io?(socket)
 
-      response.not_nil!.body.chomp.should eq "pong"
+      response.not_nil!.body.should eq "pong"
     ensure
       server.try &.close
     end

diff --git a/src/socks.cr b/src/socks.cr
@@ -71,6 +71,18 @@ class Socks < IPSocket
     end
   end
 
+  def self.udp(*args)
+    UDP.new(*args)
+  end
+
+  def self.client(*args)
+    Client.new(*args)
+  end
+
+  def self.server(*args)
+    Server.new(*args)
+  end
+
   def main_connect(addr : String, port : Int)
     connect_host
     connect_remote(addr, port)
@@ -161,3 +173,6 @@ require "./connection_request.cr"
 require "./connection_response.cr"
 require "./request.cr"
 require "./reply.cr"
+
+require "./socks/client.cr"
+require "./socks/udp.cr"
diff --git a/src/socks/client.cr b/src/socks/client.cr
@@ -8,17 +8,12 @@ class Socks::Client < HTTP::Client
 
   def socket
     socket = @socket
-    STDOUT.puts "1"
     return socket if socket
 
     hostname = @host.starts_with?('[') && @host.ends_with?(']') ? @host[1..-2] : @host
-    STDOUT.puts "2"
-    pp! hostname, @port
     socket = Socks.new hostname, @port
-    STDOUT.puts "2.1"
     socket.read_timeout = @read_timeout if @read_timeout
     socket.sync = false
-    STDOUT.puts "3"
     @socket = socket
 
     {% if !flag?(:without_openssl) %}
@@ -27,13 +22,6 @@ class Socks::Client < HTTP::Client
         @socket = socket = tls_socket
       end
     {% end %}
-    STDOUT.puts "4"
     socket
   end
 end
-
-client = Socks::Client.new("www.example.com", host_addr: "127.0.0.1", host_port: 1080)
-response = client.get("/")
-pp! response.status_code      # => 200
-pp! response.body.lines.first # => "<!doctype html>"
-client.close
diff --git a/src/socks/udp.cr b/src/socks/udp.cr
@@ -0,0 +1,7 @@
+require "../socks.cr"
+
+class Socks::UDP < UDPSocket
+  def initialize(*args, @host_addr : String, @host_port : Int32)
+    super(*args)
+  end
+end