From ffa2bbf1e4ee60e8fb99c4a28c4a4edec9918c95 Mon Sep 17 00:00:00 2001
From: Cloud Wu <cloudwu@gmail.com>
Date: Wed, 15 Jun 2016 16:48:02 +0800
Subject: [PATCH] support mongo auth_scram_sha1

---
 lualib-src/lua-crypt.c | 20 +++++++++++
 lualib/mongo.lua       | 81 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 100 insertions(+), 1 deletion(-)

diff --git a/lualib-src/lua-crypt.c b/lualib-src/lua-crypt.c
index 5322b2d6..7910d367 100644
--- a/lualib-src/lua-crypt.c
+++ b/lualib-src/lua-crypt.c
@@ -878,6 +878,25 @@ lb64decode(lua_State *L) {
 	return 1;
 }
 
+static int
+lxor_str(lua_State *L) {
+	size_t len1,len2;
+	const char *s1 = luaL_checklstring(L,1,&len1);
+	const char *s2 = luaL_checklstring(L,2,&len2);
+	if (len2 == 0) {
+		return luaL_error(L, "Can't xor empty string");
+	}
+	luaL_Buffer b;
+	char * buffer = luaL_buffinitsize(L, &b, len1);
+	int i;
+	for (i=0;i<len1;i++) {
+		buffer[i] = s1[i] ^ s2[i % len2];
+	}
+	luaL_addsize(&b, len1);
+	luaL_pushresult(&b);
+	return 1;
+}
+
 // defined in lsha1.c
 int lsha1(lua_State *L);
 int lhmac_sha1(lua_State *L);
@@ -906,6 +925,7 @@ luaopen_crypt(lua_State *L) {
 		{ "sha1", lsha1 },
 		{ "hmac_sha1", lhmac_sha1 },
 		{ "hmac_hash", lhmac_hash },
+		{ "xor_str", lxor_str },
 		{ NULL, NULL },
 	};
 	luaL_newlib(L,l);
diff --git a/lualib/mongo.lua b/lualib/mongo.lua
index 895adb62..a8acee9e 100644
--- a/lualib/mongo.lua
+++ b/lualib/mongo.lua
@@ -4,6 +4,7 @@ local socketchannel	= require "socketchannel"
 local skynet = require "skynet"
 local driver = require "mongo.driver"
 local md5 =	require	"md5"
+local crypt = require "crypt"
 local rawget = rawget
 local assert = assert
 
@@ -86,7 +87,8 @@ local function mongo_auth(mongoc)
 
 	return function()
 		if user	~= nil and pass	~= nil then
-			assert(mongoc:auth(user, pass))
+--			assert(mongoc:auth(user, pass))
+			assert(mongoc:auth_scram_sha1(user, pass))
 		end
 		local rs_data =	mongoc:runCommand("ismaster")
 		if rs_data.ok == 1 then
@@ -184,6 +186,83 @@ function mongo_client:auth(user,password)
 	return result.ok ==	1
 end
 
+local function salt_password(password, salt, iter)
+	salt = salt .. "\0\0\0\1"
+	local output = crypt.hmac_sha1(password, salt)
+	local inter = output
+	for i=2,iter do
+		inter = crypt.hmac_sha1(password, inter)
+		output = crypt.xor_str(output, inter)
+	end
+	return output
+end
+
+function mongo_client:auth_scram_sha1(username,password)
+	local user = string.gsub(string.gsub(username, '=', '=3D'), ',' , '=2C')
+	local nonce = crypt.base64encode(crypt.randomkey())
+	local first_bare = "n="  .. user .. ",r="  .. nonce
+	local sasl_start_payload = crypt.base64encode("n,," .. first_bare)
+	local r
+
+	r = self:runCommand("saslStart",1,"autoAuthorize",1,"mechanism","SCRAM-SHA-1","payload",sasl_start_payload)
+	if r.ok ~= 1 then
+		return false
+	end
+
+	local conversationId = r['conversationId']
+	local server_first = r['payload']
+	local parsed_s = crypt.base64decode(server_first)
+	local parsed_t = {}
+	for k, v in string.gmatch(parsed_s, "(%w+)=([^,]*)") do
+		parsed_t[k] = v
+	end
+	local iterations = tonumber(parsed_t['i'])
+	local salt = parsed_t['s']
+	local rnonce = parsed_t['r']
+
+	if not string.sub(rnonce, 1, 12) == nonce then
+		skynet.error("Server returned an invalid nonce.")
+		return false
+	end
+	local without_proof = "c=biws,r=" .. rnonce
+	local pbkdf2_key = md5.sumhexa(string.format("%s:mongo:%s",username,password))
+	local salted_pass = salt_password(pbkdf2_key, crypt.base64decode(salt), iterations)
+	local client_key = crypt.hmac_sha1(salted_pass, "Client Key")
+	local stored_key = crypt.sha1(client_key)
+	local auth_msg = first_bare .. ',' .. parsed_s .. ',' .. without_proof
+	local client_sig = crypt.hmac_sha1(stored_key, auth_msg)
+	local client_key_xor_sig = crypt.xor_str(client_key, client_sig)
+	local client_proof = "p=" .. crypt.base64encode(client_key_xor_sig)
+	local client_final = crypt.base64encode(without_proof .. ',' .. client_proof)
+	local server_key = crypt.hmac_sha1(salted_pass, "Server Key")
+	local server_sig = crypt.base64encode(crypt.hmac_sha1(server_key, auth_msg))
+
+	r = self:runCommand("saslContinue",1,"conversationId",conversationId,"payload",client_final)
+	if r.ok ~= 1 then
+		return false
+	end
+	parsed_s = crypt.base64decode(r['payload'])
+	parsed_t = {}
+	for k, v in string.gmatch(parsed_s, "(%w+)=([^,]*)") do
+		parsed_t[k] = v
+	end
+	if parsed_t['v'] ~= server_sig then
+		skynet.error("Server returned an invalid signature.")
+		return false
+	end
+	if not r.done then
+		r = self:runCommand("saslContinue",1,"conversationId",conversationId,"payload","")
+		if r.ok ~= 1 then
+			return false
+		end
+		if not r.done then
+			skynet.error("SASL conversation failed to complete.")
+			return false
+		end
+	end
+	return true
+end
+
 function mongo_client:logout()
 	local result = self:runCommand "logout"
 	return result.ok ==	1