Current section
Files
Jump to
Current section
Files
src/matrix_olm_session.erl
%%%===================================================================
%%% matrix_olm_session.erl — Pure Erlang Olm inbound session
%%%
%%% Implements the Olm protocol (X3DH + Double Ratchet):
%%% https://gitlab.matrix.org/matrix-org/olm/-/blob/master/docs/olm.md
%%%===================================================================
-module(matrix_olm_session).
-export([
create_account/0,
account_identity_keys/1,
account_ed25519_keypair/1,
account_one_time_keys/1,
account_generate_otks/2,
account_mark_otks_published/1,
account_remove_otk/2,
create_inbound/3,
decrypt/3,
pickle_account/1,
unpickle_account/1,
pickle_session/1,
unpickle_session/1
]).
-type curve25519_keypair() :: {binary(), binary()}.
-record(account, {
identity_keypair :: curve25519_keypair(),
ed25519_keypair :: {binary(), binary()},
one_time_keys :: #{binary() => curve25519_keypair()},
published_otk_ids :: [binary()]
}).
-record(olm_session, {
root_key :: binary(),
recv_chain :: binary(),
dh_keypair :: curve25519_keypair(),
their_ratchet_pub :: binary(),
skipped = #{} :: #{}
}).
-define(OLM_ROOT_INFO, <<"OLM_ROOT">>).
-define(OLM_RATCHET_INFO, <<"OLM_RATCHET">>).
-define(OLM_KEYS_INFO, <<"OLM_KEYS">>).
-define(MSG_KEY_SEED, <<1>>).
-define(CHAIN_KEY_SEED, <<2>>).
-define(MSG_VERSION, 3).
%%%===================================================================
%%% Account Management
%%%===================================================================
-spec create_account() -> {ok, #account{}}.
create_account() ->
{ok, #account{
identity_keypair = generate_curve25519_keypair(),
ed25519_keypair = generate_ed25519_keypair(),
one_time_keys = #{},
published_otk_ids = []
}}.
-spec account_identity_keys(#account{}) -> #{binary() => binary()}.
account_identity_keys(#account{identity_keypair = {Pub, _},
ed25519_keypair = {EdPub, _}}) ->
#{<<"curve25519">> => base64url(Pub),
<<"ed25519">> => base64url(EdPub)}.
-spec account_ed25519_keypair(#account{}) -> {binary(), binary()}.
account_ed25519_keypair(#account{ed25519_keypair = KP}) -> KP.
-spec account_one_time_keys(#account{}) -> #{binary() => binary()}.
account_one_time_keys(#account{one_time_keys = OTKs, published_otk_ids = Published}) ->
maps:fold(fun(KeyId, {Pub, _}, Acc) ->
case lists:member(KeyId, Published) of
true -> Acc;
false -> maps:put(<<"curve25519:", KeyId/binary>>, base64url(Pub), Acc)
end
end, #{}, OTKs).
-spec account_generate_otks(#account{}, pos_integer()) -> #account{}.
account_generate_otks(Acc = #account{one_time_keys = OTKs}, N) ->
NewOTKs = lists:foldl(fun(_, Map) ->
maps:put(generate_key_id(), generate_curve25519_keypair(), Map)
end, OTKs, lists:seq(1, N)),
Acc#account{one_time_keys = NewOTKs}.
-spec account_mark_otks_published(#account{}) -> #account{}.
account_mark_otks_published(Acc = #account{one_time_keys = OTKs,
published_otk_ids = Published}) ->
Acc#account{published_otk_ids = lists:usort(Published ++ maps:keys(OTKs))}.
-spec account_remove_otk(#account{}, binary()) -> #account{}.
account_remove_otk(Acc = #account{one_time_keys = OTKs}, OtkPubB64) ->
Acc#account{one_time_keys = maps:filter(
fun(_, {Pub, _}) -> base64url(Pub) =/= OtkPubB64 end, OTKs)}.
%%%===================================================================
%%% Session Creation (Inbound)
%%%===================================================================
%% FIX: spec now matches implementation (4-tuple return)
-spec create_inbound(#account{}, binary(), binary()) ->
{ok, binary(), #olm_session{}, #account{}} | {error, term()}.
create_inbound(Account, SenderIdentityKeyB64, PreKeyMsgBin) ->
try
#{base_key := BaseKeyB64,
one_time_key := OtkKeyB64,
message := InnerMsg} = parse_prekey_message(PreKeyMsgBin),
SenderIdentityKey = b64_decode(SenderIdentityKeyB64),
BaseKey = b64_decode(BaseKeyB64),
OtkPub = b64_decode(OtkKeyB64),
#account{identity_keypair = {_IdPub, IdPriv},
one_time_keys = OTKs} = Account,
{_, OtkPrivKey} = find_otk_by_pub(OTKs, OtkPub),
%% X3DH: 3 ECDH operations
S1 = ecdh(IdPriv, BaseKey),
S2 = ecdh(OtkPrivKey, SenderIdentityKey),
S3 = ecdh(OtkPrivKey, BaseKey),
%% Derive root key and chain key
IKM = <<(binary:copy(<<16#FF>>, 32))/binary, S1/binary, S2/binary, S3/binary>>,
<<RootKey:32/binary, ChainKey:32/binary>> =
crypto:hkdf(sha256, IKM, <<0:256>>, ?OLM_ROOT_INFO, 64),
Session = #olm_session{
root_key = RootKey,
recv_chain = ChainKey,
dh_keypair = generate_curve25519_keypair(),
their_ratchet_pub = BaseKey
},
Account2 = account_remove_otk(Account, OtkKeyB64),
{ok, Plaintext, Session2} = decrypt_ratchet(Session, InnerMsg, BaseKey),
{ok, Plaintext, Session2, Account2}
catch C:R:St ->
{error, {inbound_session_failed, C, R, St}}
end.
%%%===================================================================
%%% Message Decryption
%%%===================================================================
-spec decrypt(#olm_session{}, non_neg_integer(), binary()) ->
{ok, binary(), #olm_session{}} | {error, term()}.
decrypt(Session, 1, CiphertextBin) ->
try
#{ratchet_key := RatchetKeyB64,
message := Ciphertext} = parse_ratchet_message(CiphertextBin),
decrypt_ratchet(Session, Ciphertext, b64_decode(RatchetKeyB64))
catch C:R ->
{error, {decrypt_failed, C, R}}
end;
decrypt(_Session, Type, _) ->
{error, {unsupported_msg_type, Type}}.
%%%===================================================================
%%% Internal — Double Ratchet
%%%===================================================================
decrypt_ratchet(Session = #olm_session{
their_ratchet_pub = CurrentPub,
recv_chain = ChainKey,
root_key = RootKey,
dh_keypair = {_MyPub, MyPriv}
}, Ciphertext, TheirRatchetKey) ->
{ChainKey2, Session2} =
case TheirRatchetKey =:= CurrentPub of
true ->
{ChainKey, Session};
false ->
DhSecret = ecdh(MyPriv, TheirRatchetKey),
<<NewRootKey:32/binary, NewChainKey:32/binary>> =
crypto:hkdf(sha256, <<RootKey/binary, DhSecret/binary>>,
<<0:256>>, ?OLM_RATCHET_INFO, 64),
S2 = Session#olm_session{
root_key = NewRootKey,
recv_chain = NewChainKey,
dh_keypair = generate_curve25519_keypair(),
their_ratchet_pub = TheirRatchetKey
},
{NewChainKey, S2}
end,
MsgKey = hmac256(ChainKey2, ?MSG_KEY_SEED),
NewChainKey2 = hmac256(ChainKey2, ?CHAIN_KEY_SEED),
{AesKey, MacKey, AesIv} = derive_olm_keys(MsgKey),
MsgLen = byte_size(Ciphertext),
MacStart = MsgLen - 8,
<<Body:MacStart/binary, Mac:8/binary>> = Ciphertext,
ExpMac = binary:part(crypto:mac(hmac, sha256, MacKey, Body), 0, 8),
case ExpMac =:= Mac of
false -> {error, mac_mismatch};
true ->
case keylara_aes:decrypt(Body, AesKey, AesIv) of
{ok, Plaintext} ->
{ok, Plaintext, Session2#olm_session{recv_chain = NewChainKey2}};
Err -> Err
end
end.
derive_olm_keys(MsgKey) ->
<<AesKey:32/binary, MacKey:32/binary, AesIv:16/binary>> =
crypto:hkdf(sha256, MsgKey, <<>>, ?OLM_KEYS_INFO, 80),
{AesKey, MacKey, AesIv}.
%%%===================================================================
%%% Message Parsing
%%%===================================================================
parse_prekey_message(Bin) ->
Fields = decode_protobuf(Bin),
#{one_time_key => maps:get(1, Fields),
base_key => maps:get(2, Fields),
identity_key => maps:get(3, Fields),
message => maps:get(4, Fields)}.
parse_ratchet_message(Bin) ->
<<?MSG_VERSION:8, Rest/binary>> = Bin,
Fields = decode_protobuf(Rest),
#{ratchet_key => maps:get(1, Fields),
index => varint_to_integer(maps:get(2, Fields, <<0>>)),
message => maps:get(3, Fields)}.
varint_to_integer(B) when is_binary(B) -> element(1, decode_varint(B));
varint_to_integer(N) when is_integer(N) -> N.
decode_protobuf(Bin) -> decode_protobuf(Bin, #{}).
decode_protobuf(<<>>, Acc) -> Acc;
decode_protobuf(Bin, Acc) ->
{Tag, Rest} = decode_varint(Bin),
FieldNum = Tag bsr 3,
WireType = Tag band 7,
case WireType of
0 ->
{Value, Rest2} = decode_varint(Rest),
decode_protobuf(Rest2, maps:put(FieldNum, Value, Acc));
2 ->
{Len, Rest2} = decode_varint(Rest),
<<Value:Len/binary, Rest3/binary>> = Rest2,
decode_protobuf(Rest3, maps:put(FieldNum, Value, Acc));
_ -> Acc
end.
decode_varint(Bin) -> decode_varint(Bin, 0, 0).
decode_varint(<<1:1, B:7, Rest/binary>>, Shift, Acc) ->
decode_varint(Rest, Shift + 7, Acc bor (B bsl Shift));
decode_varint(<<0:1, B:7, Rest/binary>>, Shift, Acc) ->
{Acc bor (B bsl Shift), Rest}.
%%%===================================================================
%%% Crypto helpers
%%%===================================================================
ecdh(PrivKey, PubKey) ->
crypto:compute_key(ecdh, PubKey, PrivKey, x25519).
hmac256(Key, Data) ->
crypto:mac(hmac, sha256, Key, Data).
generate_curve25519_keypair() ->
crypto:generate_key(ecdh, x25519).
generate_ed25519_keypair() ->
crypto:generate_key(eddsa, ed25519).
%% FIX: keylara fallback to crypto:strong_rand_bytes if not started
generate_key_id() ->
Bytes = try
case keylara:get_entropy_bytes(4) of
{ok, B} -> B;
_ -> crypto:strong_rand_bytes(4)
end
catch _:_ ->
crypto:strong_rand_bytes(4)
end,
binary:encode_hex(Bytes).
find_otk_by_pub(OTKs, OtkPub) ->
case maps:fold(fun(_, KP = {Pub, _}, Acc) ->
case Pub =:= OtkPub of
true -> KP;
false -> Acc
end
end, not_found, OTKs) of
not_found -> throw({error, one_time_key_not_found});
KP -> KP
end.
base64url(Bin) -> base64:encode(Bin).
b64_decode(B64) ->
Padded = case byte_size(B64) rem 4 of
0 -> B64;
N -> <<B64/binary, (binary:copy(<<"=">>, 4 - N))/binary>>
end,
base64:decode(Padded).
%%%===================================================================
%%% Pickle / Unpickle
%%%===================================================================
pickle_account(Acc) -> term_to_binary(Acc).
unpickle_account(Bin) ->
try {ok, binary_to_term(Bin, [safe])}
catch _:_ -> {error, bad_pickle}
end.
pickle_session(S) -> term_to_binary(S).
unpickle_session(Bin) ->
try {ok, binary_to_term(Bin, [safe])}
catch _:_ -> {error, bad_pickle}
end.