diff --git a/src/msgpack.erl b/src/msgpack.erl index 98af3c6..01a08cc 100644 --- a/src/msgpack.erl +++ b/src/msgpack.erl @@ -129,7 +129,15 @@ unpack_stream(Bin) -> unpack_stream(Bin, []). unpack_stream(Bin, Opts0) when is_binary(Bin) -> Opts = parse_options(Opts0), try - msgpack_unpacker:unpack_stream(Bin, Opts) + Opts1=Opts?OPTION{ + known_atoms=lists:map( + fun(X) -> + erlang:atom_to_binary(X, utf8) + end, + Opts?OPTION.known_atoms + ) + }, + msgpack_unpacker:unpack_stream(Bin, Opts1) catch throw:Exception -> {error, Exception} end; diff --git a/src/msgpack.hrl b/src/msgpack.hrl index 4806543..0f47e73 100644 --- a/src/msgpack.hrl +++ b/src/msgpack.hrl @@ -36,7 +36,7 @@ -record(options_v4, { spec = new :: new | old, allow_atom = pack :: none | pack, %% allows atom when packing - known_atoms = [] :: [atom()], + known_atoms = [] :: [atom()|binary()], unpack_str = as_list :: as_binary | as_list | as_tagged_list, validate_string = false :: boolean(), pack_str = from_list :: from_binary | from_list | from_tagged_list | none, diff --git a/src/msgpack_unpacker.erl b/src/msgpack_unpacker.erl index b999af0..3484466 100644 --- a/src/msgpack_unpacker.erl +++ b/src/msgpack_unpacker.erl @@ -237,6 +237,14 @@ unpack_str_or_raw(V, ?OPTION{spec=new, as_tagged_list -> {string, unpack_str(V)} end, Rest}. +maybe_bin(Bin, ?OPTION{known_atoms=Known}) when Known=/=[] -> + case lists:member(Bin,Known) of + true -> + erlang:binary_to_existing_atom(Bin,utf8); + false -> + Bin + end; + maybe_bin(Bin, _) -> Bin. diff --git a/test/msgpack_tests.erl b/test/msgpack_tests.erl index c703530..7b3c5b5 100644 --- a/test/msgpack_tests.erl +++ b/test/msgpack_tests.erl @@ -379,6 +379,13 @@ binary_test_() -> end} ]. +atom_test_() -> + Map1=#{atom1=>atom1,atom2=><<"binary2">>}, + Bin=msgpack:pack(Map1), + ?_assertEqual({ok,#{<<"atom1">>=><<"atom1">>, + <<"atom2">>=><<"binary2">>}}, msgpack:unpack(Bin)), + ?_assertEqual({ok,Map1}, msgpack:unpack(Bin,[{known_atoms,[atom1,atom2]}])). + -define(PCNT, 5). -define(CNT, 10000).