Implement rate limiting of add_chain

1 files changed, 36 insertions, 26 deletions

diff --git a/src/catlfish.erl b/src/catlfish.erl
index dd25a84..e48f788 100644
--- a/src/catlfish.erl
+++ b/src/catlfish.erl
@@ -113,39 +113,49 @@ get_sct(Hash, TimestampedEntry) ->
             calc_sct(TimestampedEntry)
     end.
 
--spec add_chain(binary(), [binary()], normal|precert) -> {[{_,_},...]}.
-add_chain(LeafCert, CertChain, Type) ->
-    CombinedChain = [LeafCert | CertChain],
-    EntryHash = crypto:hash(sha256, CombinedChain),
+add_to_db(Type, LeafCert, CertChain, EntryHash) ->
     EntryType = case Type of
                     normal -> x509_entry;
                     precert -> precert_entry
                 end,
+    Timestamp = plop:generate_timestamp(),
+    TSE = timestamped_entry(Timestamp, EntryType, LeafCert, CertChain),
+    MTLText = serialise(#mtl{leaf_version = v1,
+                             leaf_type = timestamped_entry,
+                             entry = TSE}),
+    MTLHash = ht:leaf_hash(MTLText),
+    ExtraData =
+        case Type of
+            normal -> CertChain;
+            precert -> [LeafCert | CertChain]
+        end,
+    LogEntry =
+        list_to_binary(
+          [encode_tls_vector(MTLText, 4),
+           encode_tls_vector(
+             encode_tls_vector(
+               list_to_binary(
+                 [encode_tls_vector(C, 3) || C <- ExtraData]),
+               3),
+             4)]),
+    ok = plop:add(LogEntry, MTLHash, EntryHash),
+    {TSE, MTLHash}.
+
+get_ratelimit_token(Type) ->
+    ratelimit:get_token(Type).
+
+-spec add_chain(binary(), [binary()], normal|precert) -> {[{_,_},...]}.
+add_chain(LeafCert, CertChain, Type) ->
+    EntryHash = crypto:hash(sha256, [LeafCert | CertChain]),
     {TimestampedEntry, Hash} =
         case plop:get(EntryHash) of
             notfound ->
-                Timestamp = plop:generate_timestamp(),
-                TSE = timestamped_entry(Timestamp, EntryType, LeafCert, CertChain),
-                MTLText = serialise(#mtl{leaf_version = v1,
-                                         leaf_type = timestamped_entry,
-                                         entry = TSE}),
-                MTLHash = ht:leaf_hash(MTLText),
-                ExtraData = 
-                    case Type of
-                        normal -> CertChain;
-                        precert -> CombinedChain
-                    end,
-                LogEntry =
-                    list_to_binary(
-                      [encode_tls_vector(MTLText, 4),
-                       encode_tls_vector(
-                         encode_tls_vector(
-                           list_to_binary(
-                             [encode_tls_vector(C, 3) || C <- ExtraData]),
-                           3),
-                         4)]),
-                ok = plop:add(LogEntry, MTLHash, EntryHash),
-                {TSE, MTLHash};
+                case get_ratelimit_token(add_chain) of
+                    ok ->
+                        add_to_db(Type, LeafCert, CertChain, EntryHash);
+                    _ ->
+                        exit({internalerror, "Rate limiting"})
+                end;
             {_Index, MTLHash, DBEntry} ->
                 {MTLText, _ExtraData} = unpack_entry(DBEntry),
                 MTL = deserialise_mtl(MTLText),