From fa3108e6e2e1480915182ac6efb61a677092ef2b Mon Sep 17 00:00:00 2001
From: Konstantinos Kallas <konstantinos.kallas@hotmail.com>
Date: Mon, 17 Jul 2017 10:42:09 +0300
Subject: [PATCH] Save acquired certificates in persistent storage

---
 include/ejabberd_acme.hrl |  4 --
 src/ejabberd_acme.erl     | 80 ++++++++++++++++++++++++++++++++++-----
 2 files changed, 71 insertions(+), 13 deletions(-)

diff --git a/include/ejabberd_acme.hrl b/include/ejabberd_acme.hrl
index 0fea5bdf5..cb711b272 100644
--- a/include/ejabberd_acme.hrl
+++ b/include/ejabberd_acme.hrl
@@ -11,10 +11,6 @@
 	  key :: jose_jwk:key()
 	 }).
 
--record(data_cert, {
-	  domain  :: list(),
-	  pem     :: binary()
-	 }).
 
 
 
diff --git a/src/ejabberd_acme.erl b/src/ejabberd_acme.erl
index 5d0608c93..527cef13f 100644
--- a/src/ejabberd_acme.erl
+++ b/src/ejabberd_acme.erl
@@ -39,6 +39,7 @@ is_valid_account_opt("old-account") -> true;
 is_valid_account_opt("new-account") -> true;
 is_valid_account_opt(_) -> false.
 
+
 %%
 %% Get Certificate
 %%
@@ -403,6 +404,10 @@ is_error(_) -> false.
 data_empty() ->
     [].
 
+%%
+%% Account
+%%
+
 data_get_account(Data) ->
     case lists:keyfind(account, 1, Data) of
 	{account, #data_acc{id = AccId, key = PrivateKey}} ->
@@ -415,6 +420,27 @@ data_set_account(Data, {AccId, PrivateKey}) ->
     NewAcc = {account, #data_acc{id = AccId, key = PrivateKey}},
     lists:keystore(account, 1, Data, NewAcc).
 
+%%
+%% Certificates
+%% 
+
+data_get_certificates(Data) ->
+    case lists:keyfind(certs, 1, Data) of
+	{certs, Certs} ->
+	    {ok, Certs};
+        false ->
+	    {ok, []}
+    end.
+
+data_set_certificates(Data, NewCerts) -> 
+    lists:keystore(certs, 1, Data, {certs, NewCerts}).
+
+%% ATM we preserve one certificate for each domain
+data_add_certificate(Data, {Domain, PemCert}) ->
+    {ok, Certs} = data_get_certificates(Data),
+    NewCerts = lists:keystore(Domain, 1, Certs, {Domain, PemCert}),
+    data_set_certificates(Data, NewCerts).
+
 
 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
 %%
@@ -475,6 +501,15 @@ read_account_persistent() ->
     {ok, Data} = read_persistent(),
     data_get_account(Data).
 
+read_certificates_persistent() ->
+    {ok, Data} = read_persistent(),
+    data_get_certificates(Data).
+
+add_certificate_persistent({Domain, PemCert}) ->
+    {ok, Data} = read_persistent(),
+    NewData = data_add_certificate(Data, {Domain, PemCert}),
+    ok = write_persistent(NewData).
+
 save_certificate({error, _, _} = Error) ->
     Error;
 save_certificate({ok, DomainName, Cert}) ->
@@ -482,22 +517,30 @@ save_certificate({ok, DomainName, Cert}) ->
 	{ok, CertDir} = get_config_cert_dir(),
 	DomainString = bitstring_to_list(DomainName),
 	CertificateFile = filename:join([CertDir, DomainString ++ "_cert.pem"]),
-	case file:write_file(CertificateFile, Cert) of
-	    ok ->
-		{ok, DomainName, saved};
-	    {error, Reason} ->
-		?ERROR_MSG("Error: ~p saving certificate at file: ~p",
-			   [Reason, CertificateFile]),
-	        throw({error, DomainName, saving})
-	end
+	%% TODO: At some point do the following using a Transaction so
+	%% that there is no certificate saved if it cannot be added in
+	%% certificate persistent storage
+	write_cert(CertificateFile, Cert, DomainName),
+	add_certificate_persistent({DomainName, Cert}),
+	{ok, DomainName, saved}
     catch
 	throw:Throw ->
 	    Throw;
 	E:R ->
-	    ?ERROR_MSG("unknown ~p:~p", [E,R]),
+	    ?ERROR_MSG("Unknown ~p:~p, ~p", [E, R, erlang:get_stacktrace()]), 
 	    {error, DomainName, saving}
     end.
 
+write_cert(CertificateFile, Cert, DomainName) ->
+    case file:write_file(CertificateFile, Cert) of
+	ok ->
+	    {ok, DomainName, saved};
+	{error, Reason} ->
+	    ?ERROR_MSG("Error: ~p saving certificate at file: ~p",
+		       [Reason, CertificateFile]),
+	    throw({error, DomainName, saving})
+    end.
+
 get_config_acme() ->
     case ejabberd_config:get_option(acme, undefined) of
 	undefined ->
@@ -535,6 +578,25 @@ get_config_cert_dir() ->
 	    {ok, CertDir}
     end.
 
+%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+%%
+%% Transaction Fun
+%%
+%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+
+transaction([{Fun, Rollback} | Rest]) ->
+    try
+        {ok, Result} = Fun(),
+        [Result | transaction(Rest)]
+    catch Type:Reason ->
+        Rollback(),
+        erlang:raise(Type, Reason, erlang:get_stacktrace())
+    end;
+transaction([Fun | Rest]) ->
+    % not every action require cleanup on error
+    transaction([{Fun, fun () -> ok end} | Rest]);
+transaction([]) -> [].
+
 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
 %%
 %% Debugging Funcs -- They are only used for the development phase