erlang_adk

%% @doc Short-lived worker that resolves an opaque credential reference and
%% invokes an authentication provider.
%%
%% The child specification contains no secret material. Raw credentials exist
%% only on this process' callback stack and in its linked provider process;
%% neither is retained in gen_server state or emitted as an error.
-module(adk_token_refresh_worker).

-behaviour(gen_server).

-export([start_link/2, perform/2]).
-export([init/1, handle_call/3, handle_cast/2, handle_info/2,
         terminate/2, code_change/3, format_status/1]).

-record(state, {
    manager :: pid(),
    generation :: reference(),
    manager_monitor :: reference(),
    provider_pid = undefined :: undefined | pid(),
    provider_alias = undefined :: undefined | reference(),
    manager_alias = undefined :: undefined | reference(),
    deadline_ms = undefined :: undefined | integer()
}).

-define(MAX_WORK_BYTES, 262144).
-define(MAX_CONTEXT_BYTES, 131072).
-define(MAX_CREDENTIAL_BYTES, 1048576).
-define(MAX_PROVIDER_RESULT_BYTES, 1048576).
-define(MAX_REFRESH_TOKEN_BYTES, 65536).
-define(MAX_ID_BYTES, 4096).
-define(WORKER_MAX_HEAP_WORDS, 500000).
-define(PROVIDER_MAX_HEAP_WORDS, 300000).

-spec start_link(pid(), reference()) -> gen_server:start_ret().
start_link(Manager, Generation)
  when is_pid(Manager), is_reference(Generation) ->
    gen_server:start_link(?MODULE, {Manager, Generation}, []).

-spec perform(pid(), map()) -> ok.
perform(Worker, Work) when is_pid(Worker), is_map(Work) ->
    case bounded_term(Work, ?MAX_WORK_BYTES) of
        true -> gen_server:cast(Worker, {perform, Work});
        false -> gen_server:cast(Worker, invalid_perform)
    end;
perform(Worker, _Work) when is_pid(Worker) ->
    gen_server:cast(Worker, invalid_perform).

init({Manager, Generation}) ->
    process_flag(trap_exit, true),
    _ = process_flag(message_queue_data, off_heap),
    _ = process_flag(
          max_heap_size,
          #{size => ?WORKER_MAX_HEAP_WORDS, kill => true,
            error_logger => false, include_shared_binaries => true}),
    ManagerMonitor = erlang:monitor(process, Manager),
    Manager ! {auth_refresh_ready, Generation, self()},
    {ok, #state{manager = Manager,
                generation = Generation,
                manager_monitor = ManagerMonitor}}.

handle_call(_Request, _From, State) ->
    {reply, {error, unsupported}, State}.

handle_cast({perform, Work}, State = #state{provider_pid = undefined}) ->
    case valid_work(Work) of
        true ->
            ProviderAlias = erlang:alias([explicit_unalias]),
            ManagerAlias = maps:get(manager_alias, Work),
            Deadline = maps:get(deadline_ms, Work),
            ProviderFun = fun() ->
                Result = perform_refresh(Work),
                CompletedAt = monotonic_ms(),
                _ = erlang:send(
                      ProviderAlias,
                      {ProviderAlias, provider_result, self(), CompletedAt,
                       Result},
                      [noconnect, nosuspend]),
                ok
            end,
            ProviderPid = spawn_opt(
                            ProviderFun,
                            [link, {message_queue_data, off_heap},
                             {max_heap_size,
                              #{size => ?PROVIDER_MAX_HEAP_WORDS,
                                kill => true, error_logger => false,
                                include_shared_binaries => true}}]),
            {noreply, State#state{provider_pid = ProviderPid,
                                  provider_alias = ProviderAlias,
                                  manager_alias = ManagerAlias,
                                  deadline_ms = Deadline}};
        false ->
            send_invalid_work_result(Work, State),
            {stop, normal, State}
    end;
handle_cast(invalid_perform, State = #state{provider_pid = undefined}) ->
    {stop, normal, State};
handle_cast(_Request, State) ->
    {noreply, State}.

handle_info({ProviderAlias, provider_result, ProviderPid, CompletedAt, Result},
            State = #state{generation = Generation,
                           provider_pid = ProviderPid,
                           provider_alias = ProviderAlias,
                           manager_alias = ManagerAlias,
                           deadline_ms = Deadline}) ->
    unlink(ProviderPid),
    _ = safe_unalias(ProviderAlias),
    SafeResult = case CompletedAt =< Deadline of
        false -> {error, refresh_timeout};
        true ->
            case bounded_term(Result, ?MAX_PROVIDER_RESULT_BYTES) of
                true -> Result;
                false -> {error, invalid_provider_response}
            end
    end,
    _ = send_manager_result(ManagerAlias, Generation, CompletedAt, SafeResult),
    {stop, normal,
     State#state{provider_pid = undefined, provider_alias = undefined}};
handle_info({'EXIT', ProviderPid, _Reason},
            State = #state{generation = Generation,
                           provider_pid = ProviderPid,
                           provider_alias = ProviderAlias,
                           manager_alias = ManagerAlias}) ->
    _ = safe_unalias(ProviderAlias),
    _ = send_manager_result(
          ManagerAlias, Generation, monotonic_ms(),
          {error, provider_process_failed}),
    {stop, normal,
     State#state{provider_pid = undefined, provider_alias = undefined}};
handle_info({'DOWN', Monitor, process, _Manager, _Reason},
            State = #state{manager_monitor = Monitor}) ->
    clear_provider(State),
    {stop, normal,
     State#state{provider_pid = undefined, provider_alias = undefined}};
handle_info(_Info, State) ->
    {noreply, State}.

terminate(_Reason, State) ->
    clear_provider(State),
    ok.

code_change(_OldVersion, State, _Extra) ->
    {ok, State}.

%% A cast can transiently contain provider context and a result can contain an
%% access token. Suppress both from status/error logging.
format_status(Status) ->
    maps:map(
      fun(state, #state{manager = Manager,
                        generation = Generation,
                        provider_pid = ProviderPid,
                        deadline_ms = Deadline}) ->
              #{manager => Manager,
                generation => Generation,
                provider_running => is_pid(ProviderPid),
                deadline_ms => Deadline};
         (message, _Message) -> adk_secret_redactor:marker();
         (log, _Log) -> [];
         (reason, _Reason) -> adk_secret_redactor:marker();
         (_Key, Value) -> adk_secret_redactor:redact(Value)
      end, Status).

perform_refresh(#{store_module := StoreModule,
                  store_handle := StoreHandle,
                  principal := Principal,
                  provider := Provider,
                  credential_ref := CredentialRef,
                  provider_module := ProviderModule,
                  context := Context,
                  deadline_ms := Deadline}) ->
    case within_deadline(Deadline) of
        false ->
            {error, refresh_timeout};
        true ->
            case fetch_credential(StoreModule, StoreHandle, Principal,
                                  Provider, CredentialRef) of
                {ok, Credential} ->
                    case within_deadline(Deadline) of
                        true ->
                            invoke_provider(
                              ProviderModule, Credential, Context,
                              StoreModule, StoreHandle, Principal, Provider,
                              CredentialRef, Deadline);
                        false ->
                            {error, refresh_timeout}
                    end;
                {error, _} = Error ->
                    Error
            end
    end;
perform_refresh(_InvalidWork) ->
    {error, invalid_refresh_work}.

fetch_credential(StoreModule, StoreHandle, Principal, Provider,
                 CredentialRef) ->
    try StoreModule:fetch(StoreHandle, Principal, Provider, CredentialRef) of
        {ok, Credential} when is_map(Credential) ->
            case bounded_term(Credential, ?MAX_CREDENTIAL_BYTES) of
                true -> {ok, Credential};
                false -> {error, invalid_provider_response}
            end;
        {error, not_found} -> {error, credential_not_found};
        {error, _Reason} -> {error, credential_store_unavailable};
        _Other -> {error, credential_store_unavailable}
    catch
        _Class:_Reason -> {error, credential_store_unavailable}
    end.

invoke_provider(ProviderModule, Credential, Context,
                StoreModule, StoreHandle, Principal, Provider,
                CredentialRef, Deadline) ->
    Seeds = safe_seed_values(Credential),
    Rotator = credential_rotator(StoreModule, StoreHandle, Principal,
                                 Provider, CredentialRef, Credential,
                                 Deadline),
    ProviderContext = Context#{credential_rotator => Rotator},
    try ProviderModule:refresh(Credential, ProviderContext) of
        {ok, Token} = Success ->
            case bounded_term(Success, ?MAX_PROVIDER_RESULT_BYTES) of
                true -> {ok, Token};
                false -> {error, invalid_provider_response}
            end;
        {error, credential_rotation_conflict} ->
            {error, credential_rotation_conflict};
        {error, credential_rotation_failed} ->
            {error, credential_rotation_failed};
        {error, Reason} = Failure ->
            case bounded_term(Failure, ?MAX_PROVIDER_RESULT_BYTES) of
                true ->
                    {error, {provider_error,
                             safe_redact(Reason, Seeds)}};
                false ->
                    {error, invalid_provider_response}
            end;
        _Other ->
            {error, invalid_provider_response}
    catch
        Class:Reason ->
            case bounded_term({Class, Reason},
                              ?MAX_PROVIDER_RESULT_BYTES) of
                true ->
                    {error, {provider_exception, Class,
                             safe_redact(Reason, Seeds)}};
                false ->
                    {error, invalid_provider_response}
            end
    end.

credential_rotator(StoreModule, StoreHandle, Principal, Provider,
                   CredentialRef, OriginalCredential, Deadline) ->
    fun(ExpectedCredential, NewRefreshToken) ->
        case monotonic_ms() =< Deadline of
            false ->
                {error, deadline_exceeded};
            true ->
                case ExpectedCredential =:= OriginalCredential of
                    true ->
                        rotate_refresh_token(
                          StoreModule, StoreHandle, Principal, Provider,
                          CredentialRef, ExpectedCredential,
                          NewRefreshToken);
                    false ->
                        {error, invalid_credential}
                end
        end
    end.

rotate_refresh_token(_StoreModule, _StoreHandle, _Principal, _Provider,
                     _CredentialRef,
                     #{refresh_token := CurrentRefreshToken},
                     CurrentRefreshToken)
  when is_binary(CurrentRefreshToken), byte_size(CurrentRefreshToken) > 0,
       byte_size(CurrentRefreshToken) =< ?MAX_REFRESH_TOKEN_BYTES ->
    ok;
rotate_refresh_token(StoreModule, StoreHandle, Principal, Provider,
                     CredentialRef,
                     #{refresh_token := CurrentRefreshToken} = Expected,
                     NewRefreshToken)
  when is_binary(CurrentRefreshToken), byte_size(CurrentRefreshToken) > 0,
       byte_size(CurrentRefreshToken) =< ?MAX_REFRESH_TOKEN_BYTES,
       is_binary(NewRefreshToken), byte_size(NewRefreshToken) > 0,
       byte_size(NewRefreshToken) =< ?MAX_REFRESH_TOKEN_BYTES,
       is_map(Expected) ->
    Replacement = Expected#{refresh_token => NewRefreshToken},
    case bounded_term(Expected, ?MAX_CREDENTIAL_BYTES) andalso
         bounded_term(Replacement, ?MAX_CREDENTIAL_BYTES) of
        true ->
            safe_compare_and_swap(StoreModule, StoreHandle, Principal,
                                  Provider, CredentialRef, Expected,
                                  Replacement);
        false ->
            {error, invalid_credential}
    end;
rotate_refresh_token(_StoreModule, _StoreHandle, _Principal, _Provider,
                     _CredentialRef, _Expected, _NewRefreshToken) ->
    {error, invalid_credential}.

safe_compare_and_swap(StoreModule, StoreHandle, Principal, Provider,
                      CredentialRef, Expected, Replacement) ->
    try StoreModule:compare_and_swap(
          StoreHandle, Principal, Provider, CredentialRef,
          Expected, Replacement) of
        ok -> ok;
        {error, conflict} -> {error, conflict};
        {error, _Reason} -> {error, unavailable};
        _Other -> {error, unavailable}
    catch
        _:_ -> {error, unavailable}
    end.

safe_seed_values(Credential) ->
    try adk_secret_redactor:seed_values(Credential)
    catch _:_ -> []
    end.

safe_redact(Reason, Seeds) ->
    try adk_secret_redactor:redact(Reason, Seeds)
    catch _:_ -> adk_secret_redactor:marker()
    end.

stop_provider(undefined) -> ok;
stop_provider(ProviderPid) when is_pid(ProviderPid) ->
    unlink(ProviderPid),
    exit(ProviderPid, kill),
    ok.

clear_provider(#state{provider_pid = ProviderPid,
                      provider_alias = ProviderAlias}) ->
    _ = safe_unalias(ProviderAlias),
    stop_provider(ProviderPid).

safe_unalias(undefined) -> ok;
safe_unalias(Alias) when is_reference(Alias) ->
    _ = catch erlang:unalias(Alias),
    ok.

send_manager_result(ManagerAlias, Generation, CompletedAt, Result)
  when is_reference(ManagerAlias), is_reference(Generation),
       is_integer(CompletedAt) ->
    _ = catch erlang:send(
                ManagerAlias,
                {ManagerAlias, auth_refresh_result, Generation, self(),
                 CompletedAt, Result},
                [noconnect, nosuspend]),
    ok.

send_invalid_work_result(Work, #state{generation = Generation}) ->
    case {maps:find(manager_alias, Work), maps:find(deadline_ms, Work)} of
        {{ok, ManagerAlias}, {ok, _Deadline}}
          when is_reference(ManagerAlias) ->
            send_manager_result(
              ManagerAlias, Generation, monotonic_ms(),
              {error, invalid_refresh_work});
        _ ->
            ok
    end,
    ok.

valid_work(#{store_module := StoreModule,
             store_handle := StoreHandle,
             principal := Principal,
             provider := Provider,
             credential_ref := CredentialRef,
             provider_module := ProviderModule,
             context := Context,
             deadline_ms := Deadline,
             manager_alias := ManagerAlias} = Work) ->
    lists:sort(maps:keys(Work)) =:=
        [context, credential_ref, deadline_ms, manager_alias, principal,
         provider, provider_module, store_handle, store_module] andalso
    bounded_term(Work, ?MAX_WORK_BYTES) andalso
    is_atom(StoreModule) andalso StoreModule =/= undefined andalso
    valid_server(StoreHandle) andalso
    valid_identity(Principal) andalso valid_identity(Provider) andalso
    adk_credential_store:is_ref(CredentialRef) andalso
    is_atom(ProviderModule) andalso ProviderModule =/= undefined andalso
    is_map(Context) andalso bounded_term(Context, ?MAX_CONTEXT_BYTES) andalso
    is_integer(Deadline) andalso is_reference(ManagerAlias);
valid_work(_Work) -> false.

valid_server(Value) when is_pid(Value) -> true;
valid_server(Value) when is_atom(Value) -> Value =/= undefined;
valid_server(_Value) -> false.

valid_identity(Value) when is_binary(Value) ->
    byte_size(Value) > 0 andalso byte_size(Value) =< ?MAX_ID_BYTES;
valid_identity(Value) when is_atom(Value) -> Value =/= undefined;
valid_identity(_Value) -> false.

bounded_term(Term, Maximum) ->
    try erlang:external_size(Term) =< Maximum
    catch _:_ -> false
    end.

monotonic_ms() ->
    erlang:monotonic_time(millisecond).

within_deadline(Deadline) when is_integer(Deadline) ->
    monotonic_ms() =< Deadline.