From 6b6b90b390317d5e76bfc151270bb28ba0621524 Mon Sep 17 00:00:00 2001
From: Jim Menard <jim@10gen.com>
Date: Wed, 21 Jan 2009 11:26:18 -0500
Subject: [PATCH] Database authentication.

---
 lib/mongo/db.rb  | 35 +++++++++++++++++++++++++++++++++++
 tests/test_db.rb | 45 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+)

diff --git a/lib/mongo/db.rb b/lib/mongo/db.rb
index 6ef0fe4..3c4728a 100644
--- a/lib/mongo/db.rb
+++ b/lib/mongo/db.rb
@@ -15,6 +15,7 @@
 # ++
 
 require 'socket'
+require 'md5'
 require 'mutex_m'
 require 'mongo/collection'
 require 'mongo/message'
@@ -32,6 +33,7 @@ module XGen
         SYSTEM_NAMESPACE_COLLECTION = "system.namespaces"
         SYSTEM_INDEX_COLLECTION = "system.indexes"
         SYSTEM_PROFILE_COLLECTION = "system.profile"
+        SYSTEM_USER_COLLECTION = "system.users"
         SYSTEM_COMMAND_COLLECTION = "$cmd"
 
         # Strict mode enforces collection existence checks. When +true+,
@@ -119,6 +121,33 @@ module XGen
           raise "error: failed to connect to any given host:port" unless @socket
         end
 
+        # Add a new user to the database.
+        def add_user(username, password)
+          coll = collection(SYSTEM_USER_COLLECTION)
+          coll.insert(:user => username, :pwd => hash_password(password))
+        end
+
+        def delete_user(username)
+          coll = collection(SYSTEM_USER_COLLECTION)
+          coll.remove(:user => username)
+        end
+
+        # Returns true if +username+ has +password+ in
+        # +SYSTEM_USER_COLLECTION+. +name+ is username, +password+ is
+        # plaintext password.
+        def authenticate(username, password)
+          doc = db_command(:getnonce => 1)
+          raise "error retrieving nonce: #{doc}" unless ok?(doc)
+          nonce = doc['nonce']
+
+          auth = OrderedHash.new
+          auth['authenticate'] = 1
+          auth['user'] = username
+          auth['nonce'] = nonce
+          auth['key'] = MD5.md5("#{nonce}#{username}#{hash_password(password)}").to_s
+          ok?(db_command(auth))
+        end
+
         # Returns an array of collection names. Each name is of the form
         # "database_name.collection_name".
         def collection_names
@@ -393,6 +422,12 @@ module XGen
           query(Collection.new(self, SYSTEM_COMMAND_COLLECTION), q).next_object
         end
 
+        private
+
+        def hash_password(plaintext)
+          MD5.new("mongo#{plaintext}").to_s
+        end
+
       end
     end
   end
diff --git a/tests/test_db.rb b/tests/test_db.rb
index b7ca949..2b8884f 100644
--- a/tests/test_db.rb
+++ b/tests/test_db.rb
@@ -1,4 +1,5 @@
 $LOAD_PATH[0,0] = File.join(File.dirname(__FILE__), '..', 'lib')
+require 'md5'
 require 'mongo'
 require 'test/unit'
 
@@ -18,6 +19,8 @@ class DBTest < Test::Unit::TestCase
     @host = ENV['MONGO_RUBY_DRIVER_HOST'] || 'localhost'
     @port = ENV['MONGO_RUBY_DRIVER_PORT'] || Mongo::DEFAULT_PORT
     @db = Mongo.new(@host, @port).db('ruby-mongo-test')
+    @spongebob = 'spongebob'
+    @spongebob_password = 'squarepants'
   end
 
   def teardown
@@ -82,4 +85,46 @@ class DBTest < Test::Unit::TestCase
     end
   end
 
+  def test_add_user
+    coll = @db.collection('system.users')
+    coll.clear
+    begin
+      assert_equal 0, coll.count
+      @db.add_user(@spongebob, @spongebob_password)
+      assert_equal 1, coll.count
+      doc = coll.find({}, :limit => 1).next_object
+      assert_equal @spongebob, doc['user']
+      assert_equal MD5.new("mongo#{@spongebob_password}").to_s, doc['pwd']
+    ensure
+      coll.clear
+    end
+  end
+
+  def test_delete_user
+    coll = @db.collection('system.users')
+    coll.clear
+    begin
+      assert_equal 0, coll.count
+      @db.add_user(@spongebob, @spongebob_password)
+      assert_equal 1, coll.count
+      @db.delete_user(@spongebob)
+      assert_equal 0, coll.count
+    ensure
+      coll.clear
+    end
+  end
+
+  def test_authenticate
+    coll = @db.collection('system.users')
+    coll.clear
+    begin
+      @db.add_user(@spongebob, @spongebob_password)
+      assert !@db.authenticate('nobody', 'nopassword')
+      assert !@db.authenticate(@spongebob, 'squareliederhosen')
+      assert @db.authenticate(@spongebob, @spongebob_password)
+    ensure
+      coll.clear
+    end
+  end
+
 end