From f00c0dfcf075bdb7c0012649fa31f3bc40821a4e Mon Sep 17 00:00:00 2001
From: Kyle Banker <kylebanker@gmail.com>
Date: Fri, 26 Aug 2011 17:35:40 -0400
Subject: [PATCH] RUBY-316 initial SSL support

---
 lib/mongo.rb                 |  1 +
 lib/mongo/connection.rb      | 15 +++++++++++---
 lib/mongo/node.rb            |  4 ++--
 lib/mongo/util/pool.rb       |  2 +-
 lib/mongo/util/ssl_socket.rb | 38 ++++++++++++++++++++++++++++++++++++
 test/unit/node_test.rb       |  6 ++++--
 6 files changed, 58 insertions(+), 8 deletions(-)
 create mode 100644 lib/mongo/util/ssl_socket.rb

diff --git a/lib/mongo.rb b/lib/mongo.rb
index e0e44fa..96b7a2a 100644
--- a/lib/mongo.rb
+++ b/lib/mongo.rb
@@ -59,6 +59,7 @@ require 'mongo/util/core_ext'
 require 'mongo/util/pool'
 require 'mongo/util/pool_manager'
 require 'mongo/util/server_version'
+require 'mongo/util/ssl_socket'
 require 'mongo/util/uri_parser'
 
 require 'mongo/collection'
diff --git a/lib/mongo/connection.rb b/lib/mongo/connection.rb
index 0e4fe61..0682b5f 100644
--- a/lib/mongo/connection.rb
+++ b/lib/mongo/connection.rb
@@ -36,7 +36,7 @@ module Mongo
     RESPONSE_HEADER_SIZE = 20
 
     attr_reader :logger, :size, :auths, :primary, :safe, :host_to_try,
-      :pool_size, :connect_timeout, :primary_pool
+      :pool_size, :connect_timeout, :primary_pool, :socket_class
 
     # Counter for generating unique request ids.
     @@current_request_id = 0
@@ -73,6 +73,7 @@ module Mongo
     #   Disabled by default.
     # @option opts [Float] :connect_timeout (nil) The number of seconds to wait before timing out a
     #   connection attempt.
+    # @option opts [Boolean] :ssl (false) If true, create the connection to the server using SSL.
     #
     # @example localhost, 27017
     #   Connection.new
@@ -636,6 +637,14 @@ module Mongo
       # Default maximum BSON object size
       @max_bson_size = Mongo::DEFAULT_MAX_BSON_SIZE
 
+      # Determine whether to use SSL.
+      @ssl = opts.fetch(:ssl, false)
+      if @ssl
+        @socket_class = Mongo::SSLSocket
+      else
+        @socket_class = ::TCPSocket
+      end
+
       # Authentication objects
       @auths = opts.fetch(:auths, [])
 
@@ -729,11 +738,11 @@ module Mongo
 
         if @connect_timeout
           Mongo::TimeoutHandler.timeout(@connect_timeout, OperationTimeout) do
-            socket = TCPSocket.new(host, port)
+            socket = @socket_class.new(host, port)
             socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
           end
         else
-          socket = TCPSocket.new(host, port)
+          socket = @socket_class.new(host, port)
           socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
         end
 
diff --git a/lib/mongo/node.rb b/lib/mongo/node.rb
index 259e70b..483847c 100644
--- a/lib/mongo/node.rb
+++ b/lib/mongo/node.rb
@@ -31,10 +31,10 @@ module Mongo
         socket = nil
         if self.connection.connect_timeout
           Mongo::TimeoutHandler.timeout(self.connection.connect_timeout, OperationTimeout) do
-            socket = TCPSocket.new(self.host, self.port)
+            socket = self.connection.socket_class.new(self.host, self.port)
           end
         else
-          socket = TCPSocket.new(self.host, self.port)
+          socket = self.connection.socket_class.new(self.host, self.port)
         end
 
         if socket.nil?
diff --git a/lib/mongo/util/pool.rb b/lib/mongo/util/pool.rb
index 2d1ce13..6c4fc39 100644
--- a/lib/mongo/util/pool.rb
+++ b/lib/mongo/util/pool.rb
@@ -110,7 +110,7 @@ module Mongo
     # therefore, it runs within a mutex.
     def checkout_new_socket
       begin
-        socket = TCPSocket.new(@host, @port)
+        socket = self.connection.socket_class.new(@host, @port)
         socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
       rescue => ex
         socket.close if socket
diff --git a/lib/mongo/util/ssl_socket.rb b/lib/mongo/util/ssl_socket.rb
new file mode 100644
index 0000000..2a4bcbe
--- /dev/null
+++ b/lib/mongo/util/ssl_socket.rb
@@ -0,0 +1,38 @@
+require 'openssl'
+
+module Mongo
+
+  # A basic wrapper over Ruby's SSLSocket that initiates
+  # a TCP connection over SSL and then provides an basic interface
+  # mirroring Ruby's TCPSocket, vis., TCPSocket#send and TCPSocket#read.
+  class SSLSocket
+
+    def initialize(host, port)
+      @socket = ::TCPSocket.new(host, port)
+      @ssl = OpenSSL::SSL::SSLSocket.new(@socket)
+      @ssl.sync_close = true
+      @ssl.connect
+    end
+
+    def setsockopt(key, value, n)
+      @socket.setsockopt(key, value, n)
+    end
+
+    # Write to the SSL socket.
+    #
+    # @param buffer a buffer to send.
+    # @param flags socket flags. Because Ruby's SSL
+    def send(buffer, flags=0)
+      @ssl.syswrite(buffer)
+    end
+
+    def read(length, buffer)
+      @ssl.sysread(length, buffer)
+    end
+
+    def close
+      @ssl.close
+    end
+
+  end
+end
diff --git a/test/unit/node_test.rb b/test/unit/node_test.rb
index d311f54..777e307 100644
--- a/test/unit/node_test.rb
+++ b/test/unit/node_test.rb
@@ -7,8 +7,10 @@ class NodeTest < Test::Unit::TestCase
   end
 
   should "refuse to connect to node without 'hosts' key" do
+    tcp = mock()
     node = Node.new(@connection, ['localhost', 27017])
-    TCPSocket.stubs(:new).returns(new_mock_socket)
+    tcp.stubs(:new).returns(new_mock_socket)
+    @connection.stubs(:socket_class).returns(tcp)
 
     admin_db = new_mock_db
     admin_db.stubs(:command).returns({'ok' => 1, 'ismaster' => 1})
@@ -17,7 +19,7 @@ class NodeTest < Test::Unit::TestCase
     @connection.expects(:log)
 
     assert node.connect
-    assert node.set_config
+    node.set_config
   end
 
   should "load a node from an array" do