From 2d103b4ae1db89a22b0a90d647632417e1b298dd Mon Sep 17 00:00:00 2001
From: HAMANO Tsukasa <hamano@cuspy.org>
Date: Sat, 19 Mar 2016 01:41:14 +0900
Subject: [PATCH] support riak authentication

---
 src/ejabberd_riak.erl     | 12 +++++------
 src/ejabberd_riak_sup.erl | 42 +++++++++++++++++++++++++++++++++++++--
 2 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/src/ejabberd_riak.erl b/src/ejabberd_riak.erl
index 22da9d2e9..90bdb4c06 100644
--- a/src/ejabberd_riak.erl
+++ b/src/ejabberd_riak.erl
@@ -28,7 +28,7 @@
 -behaviour(gen_server).
 
 %% API
--export([start_link/4, get_proc/1, make_bucket/1, put/2, put/3,
+-export([start_link/5, get_proc/1, make_bucket/1, put/2, put/3,
          get/2, get/3, get_by_index/4, delete/1, delete/2,
          count_by_index/3, get_by_index_range/5,
          get_keys/1, get_keys_by_index/3, is_connected/0,
@@ -68,8 +68,8 @@
 %%% API
 %%%===================================================================
 %% @private
-start_link(Num, Server, Port, _StartInterval) ->
-    gen_server:start_link({local, get_proc(Num)}, ?MODULE, [Server, Port], []).
+start_link(Num, Server, Port, _StartInterval, Options) ->
+    gen_server:start_link({local, get_proc(Num)}, ?MODULE, [Server, Port, Options], []).
 
 %% @private
 is_connected() ->
@@ -429,10 +429,8 @@ map_key(Obj, _, _) ->
 %%% gen_server API
 %%%===================================================================
 %% @private
-init([Server, Port]) ->
-    case riakc_pb_socket:start(
-           Server, Port,
-           [auto_reconnect]) of
+init([Server, Port, Options]) ->
+    case riakc_pb_socket:start(Server, Port, Options) of
         {ok, Pid} ->
             erlang:monitor(process, Pid),
             {ok, #state{pid = Pid}};
diff --git a/src/ejabberd_riak_sup.erl b/src/ejabberd_riak_sup.erl
index 319d1707d..7294e0055 100644
--- a/src/ejabberd_riak_sup.erl
+++ b/src/ejabberd_riak_sup.erl
@@ -103,12 +103,26 @@ init([]) ->
     StartInterval = get_start_interval(),
     Server = get_riak_server(),
     Port = get_riak_port(),
+    CACertFile = get_riak_cacertfile(),
+    Username = get_riak_username(),
+    Password = get_riak_password(),
+    Options = lists:filter(
+		fun(X) -> X /= nil end,
+		[auto_reconnect,
+		 if CACertFile /= nil -> {cacertfile ,CACertFile};
+		    true -> nil
+		 end,
+		 if (Username /= nil) and (Password /= nil) ->
+			 {credentials, Username, Password};
+		    true -> nil
+		 end
+		]),
     {ok, {{one_for_one, PoolSize*10, 1},
 	  lists:map(
 	    fun(I) ->
 		    {ejabberd_riak:get_proc(I),
 		     {ejabberd_riak, start_link,
-                      [I, Server, Port, StartInterval*1000]},
+                      [I, Server, Port, StartInterval*1000, Options]},
 		     transient, 2000, worker, [?MODULE]}
 	    end, lists:seq(1, PoolSize))}}.
 
@@ -131,6 +145,27 @@ get_riak_server() ->
 	      binary_to_list(iolist_to_binary(S))
       end, ?DEFAULT_RIAK_HOST).
 
+get_riak_cacertfile() ->
+    ejabberd_config:get_option(
+      riak_cacertfile,
+      fun(S) ->
+	      binary_to_list(iolist_to_binary(S))
+      end, nil).
+
+get_riak_username() ->
+    ejabberd_config:get_option(
+      riak_username,
+      fun(S) ->
+	      binary_to_list(iolist_to_binary(S))
+      end, nil).
+
+get_riak_password() ->
+    ejabberd_config:get_option(
+      riak_password,
+      fun(S) ->
+	      binary_to_list(iolist_to_binary(S))
+      end, nil).
+
 get_riak_port() ->
     ejabberd_config:get_option(
       riak_port,
@@ -162,6 +197,9 @@ opt_type(riak_port) -> fun (_) -> true end;
 opt_type(riak_server) -> fun (_) -> true end;
 opt_type(riak_start_interval) ->
     fun (N) when is_integer(N), N >= 1 -> N end;
+opt_type(riak_cacertfile) -> fun iolist_to_binary/1;
+opt_type(riak_username) -> fun iolist_to_binary/1;
+opt_type(riak_password) -> fun iolist_to_binary/1;
 opt_type(_) ->
     [modules, riak_pool_size, riak_port, riak_server,
-     riak_start_interval].
+     riak_start_interval, riak_cacertfile, riak_username, riak_password].