iot_cloud/src/mnesia/efka_client_store.erl

%%%-------------------------------------------------------------------
%%% @doc efka_client table store based on mnesia.
%%%-------------------------------------------------------------------
-module(efka_client_store).

-include("iot_tables.hrl").

-export([ensure_table/0, register/2, get/1, update/1, delete/1, list/0]).
-export([auth/2, verify_heartbeat/4]).

-define(TOKEN_HASH_ITERATIONS, 32).
-define(TOKEN_HASH_BYTES, 32).
-define(HEARTBEAT_TIMESTAMP_WINDOW, 120).

-spec ensure_table() -> ok | {error, term()}.
ensure_table() ->
    case mnesia:create_table(efka_client, [
        {attributes, record_info(fields, efka_client)},
        {type, set},
        {record_name, efka_client},
        {disc_copies, [node()]}
    ]) of
        {atomic, ok} ->
            ok;
        {aborted, {already_exists, efka_client}} ->
            ok;
        {aborted, Reason} ->
            {error, Reason}
    end.

-spec auth(UUID :: binary(), Token :: binary()) -> boolean().
auth(UUID, Token) when is_binary(UUID), UUID =/= <<>>, is_binary(Token) ->
    case mnesia:dirty_read(efka_client, UUID) of
        [] ->
            false;
        [#efka_client{token_hash = TokenHash, salt = Salt}] ->
            hash_token(Token, Salt) =:= TokenHash
    end.

-spec verify_heartbeat(UUID :: binary(), Timestamp :: integer(), Payload :: binary(), Mac :: binary()) -> boolean().
verify_heartbeat(UUID, Timestamp, Payload, Mac)
  when is_binary(UUID), UUID =/= <<>>, is_integer(Timestamp), is_binary(Payload), is_binary(Mac) ->
    case valid_heartbeat_timestamp(Timestamp) of
        true ->
            case mnesia:dirty_read(efka_client, UUID) of
                [] ->
                    false;
                [#efka_client{heartbeat_secret = HeartbeatSecret}] ->
                    ExpectedMac = crypto:mac(hmac, sha256, HeartbeatSecret, Payload),
                    ExpectedMac =:= Mac
            end;
        false ->
            false
    end.

-spec register(UUID :: binary(), Token :: binary()) -> ok | {error, term()}.
register(UUID, Token) when is_binary(UUID), UUID =/= <<>>, is_binary(Token), byte_size(Token) >= 32 ->
    case mnesia:transaction(fun() ->
        case mnesia:read(efka_client, UUID, write) of
            [] ->
                %% 1. 生成 16 字节随机盐
                Salt = crypto:strong_rand_bytes(16),
                TokenHash = hash_token(Token, Salt),
                HeartbeatSecret = heartbeat_secret(Token),
                ok = mnesia:write(#efka_client{
                    uuid = UUID,
                    token_hash = TokenHash,
                    salt = Salt,
                    heartbeat_secret = HeartbeatSecret,
                    timestamp = iot_util:timestamp()
                }),
                ok;
            [_] ->
                {error, already_exists}
        end
                            end) of
        {atomic, ok} ->
            ok;
        {atomic, {error, Reason}} ->
            {error, Reason};
        {aborted, Reason} ->
            {error, Reason}
    end.

-spec get(binary()) -> {ok, #efka_client{}} | not_found | {error, term()}.
get(UUID) when is_binary(UUID), UUID =/= <<>> ->
    case mnesia:transaction(fun() ->
        case mnesia:read(efka_client, UUID, read) of
            [Client = #efka_client{}] ->
                {ok, Client};
            [] ->
                not_found
        end
    end) of
        {atomic, Result} ->
            Result;
        {aborted, Reason} ->
            {error, Reason}
    end.

-spec update(#efka_client{}) -> ok | {error, term()}.
update(Client = #efka_client{uuid = UUID}) when is_binary(UUID), UUID =/= <<>> ->
    case mnesia:transaction(fun() ->
        case mnesia:read(efka_client, UUID, write) of
            [_] ->
                ok = mnesia:write(Client),
                ok;
            [] ->
                {error, not_found}
        end
    end) of
        {atomic, ok} ->
            ok;
        {atomic, {error, Reason}} ->
            {error, Reason};
        {aborted, Reason} ->
            {error, Reason}
    end.

-spec delete(binary()) -> ok | {error, term()}.
delete(UUID) when is_binary(UUID), UUID =/= <<>> ->
    case mnesia:transaction(fun() ->
        ok = mnesia:delete({efka_client, UUID}),
        ok
    end) of
        {atomic, ok} ->
            ok;
        {aborted, Reason} ->
            {error, Reason}
    end.

-spec list() -> {ok, [#efka_client{}]} | {error, term()}.
list() ->
    case mnesia:transaction(fun() ->
        mnesia:foldl(fun(Client = #efka_client{}, Acc) -> [Client | Acc] end, [], efka_client)
    end) of
        {atomic, Clients} ->
            {ok, lists:reverse(Clients)};
        {aborted, Reason} ->
            {error, Reason}
    end.

-spec hash_token(binary(), binary()) -> binary().
hash_token(Token, Salt) when is_binary(Token), is_binary(Salt) ->
    crypto:pbkdf2_hmac(sha256, Token, Salt, ?TOKEN_HASH_ITERATIONS, ?TOKEN_HASH_BYTES).

-spec heartbeat_secret(binary()) -> binary().
heartbeat_secret(Token) when is_binary(Token) ->
    crypto:hash(sha256, Token).

-spec valid_heartbeat_timestamp(integer()) -> boolean().
valid_heartbeat_timestamp(Timestamp) ->
    Now = iot_util:current_time(),
    Timestamp =< Now andalso Now - Timestamp =< ?HEARTBEAT_TIMESTAMP_WINDOW.