diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/catlfish.erl | 63 | ||||
| -rw-r--r-- | src/catlfish_sup.erl | 13 | ||||
| -rw-r--r-- | src/ratelimit.erl | 85 |
3 files changed, 134 insertions, 27 deletions
diff --git a/src/catlfish.erl b/src/catlfish.erl index dd25a84..1e2456d 100644 --- a/src/catlfish.erl +++ b/src/catlfish.erl @@ -113,39 +113,50 @@ 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) -> + CombinedChain = [LeafCert | CertChain], + EntryHash = crypto:hash(sha256, CombinedChain), {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), diff --git a/src/catlfish_sup.erl b/src/catlfish_sup.erl index 5da7b93..3defec6 100644 --- a/src/catlfish_sup.erl +++ b/src/catlfish_sup.erl @@ -9,6 +9,16 @@ start_link(_Args) -> supervisor:start_link({local, ?MODULE}, ?MODULE, []). +permanent_worker(Name, {Module, Function, Args}) -> + permanent_worker(Name, {Module, Function, Args}, [Module]). + +permanent_worker(Name, StartFunc, Modules) -> + {Name, + StartFunc, + permanent, + 10000, + worker, Modules}. + gen_http_config(Config, SSLOptions, SSLFlag) -> {ChildName, IpAddress, Port, Module} = Config, {ok, IPv4Address} = @@ -39,4 +49,5 @@ init([]) -> lager:debug("Starting servers ~p", [Servers]), {ok, {{one_for_one, 3, 10}, - Servers}}. + [permanent_worker(ratelimit, {ratelimit, start_link, []}) | + Servers]}}. diff --git a/src/ratelimit.erl b/src/ratelimit.erl new file mode 100644 index 0000000..c840cc3 --- /dev/null +++ b/src/ratelimit.erl @@ -0,0 +1,85 @@ +%%% Copyright (c) 2014-2015, NORDUnet A/S. +%%% See LICENSE for licensing information. +%%% + +-module(ratelimit). +-behaviour(gen_server). + +-export([start_link/0, stop/0]). +-export([get_token/1]). +%% gen_server callbacks. +-export([init/1, handle_call/3, terminate/2, handle_cast/2, handle_info/2, + code_change/3]). + +start_link() -> + gen_server:start_link({local, ?MODULE}, ?MODULE, + application:get_env(catlfish, ratelimits, []), []). + +stop() -> + gen_server:call(?MODULE, stop). + +get_token(Type) -> + gen_server:call(?MODULE, {get_token, Type}). + +init(Types) -> + lager:debug("starting ratelimit service"), + State = dict:from_list([{Type, {Rules, queue:new()}} || {Type, Rules} <- Types]), + {ok, State}. + +rule_interval_atom([{_, Interval}]) -> + Interval. + +rule_interval([{_, second}]) -> + 1000; +rule_interval([{_, minute}]) -> + 60*1000; +rule_interval([{_, hour}]) -> + 60*60*1000. + +rule_times([{Times, _}]) when is_integer(Times) -> + Times. + +clean_queue(Interval, Queue) -> + Now = plop:generate_timestamp(), + case queue:peek(Queue) of + {value, Item} when Item + Interval < Now -> + clean_queue(Interval, queue:drop(Queue)); + _ -> + Queue + end. + +get_token_for_type({none, Queue}) -> + {ok, {none, Queue}}; +get_token_for_type({Rules, Queue}) -> + CleanedQueue = clean_queue(rule_interval(Rules), Queue), + MaxTimes = rule_times(Rules), + QueueLength = queue:len(CleanedQueue), + if + QueueLength < MaxTimes -> + Now = plop:generate_timestamp(), + {ok, {Rules, queue:in(Now, CleanedQueue)}}; + true -> + {overload, {Rules, CleanedQueue}} + end. + +handle_call(stop, _From, State) -> + {stop, normal, stopped, State}; +handle_call({get_token, Type}, _From, State) -> + case dict:find(Type, State) of + {ok, TypeState} -> + {Result, NewTypeState} = get_token_for_type(TypeState), + {Rules, Queue} = NewTypeState, + lager:debug("current rate: ~p per ~p", [queue:len(Queue), rule_interval_atom(Rules)]), + {reply, Result, dict:store(Type, NewTypeState, State)}; + error -> + {reply, ok, State} + end. + +handle_cast(_Request, State) -> + {noreply, State}. +handle_info(_Info, State) -> + {noreply, State}. +code_change(_OldVersion, State, _Extra) -> + {ok, State}. +terminate(_Reason, _State) -> + ok. |