diff --git a/src/mtp_codec.erl b/src/mtp_codec.erl index 82a0a87..4457bf7 100644 --- a/src/mtp_codec.erl +++ b/src/mtp_codec.erl @@ -15,7 +15,7 @@ info/2, replace/4, push_back/3, is_empty/1, try_decode_packet/2, encode_packet/2, - fold_packets/4]). + fold_packets/4, fold_packets_if/4]). -export_type([codec/0]). -type state() :: any(). @@ -185,7 +185,7 @@ encode_packet(Bin, #codec{have_tls = HaveTls, end. --spec fold_packets(fun( (binary(), FoldSt, codec()) -> FoldSt ), +-spec fold_packets(fun( (binary(), FoldSt, codec()) -> {FoldSt, codec()} ), FoldSt, binary(), codec()) -> {ok, FoldSt, codec()} when @@ -199,6 +199,24 @@ fold_packets(Fun, FoldSt, Data, Codec) -> {ok, FoldSt, Codec1} end. +-spec fold_packets_if(fun( (binary(), FoldSt, codec()) -> {next | stop, FoldSt, codec()} ), + FoldSt, binary(), codec()) -> + {ok, FoldSt, codec()} + when + FoldSt :: any(). +fold_packets_if(Fun, FoldSt0, Data, Codec0) -> + case try_decode_packet(Data, Codec0) of + {ok, Decoded, Codec1} -> + case Fun(Decoded, FoldSt0, Codec1) of + {next, FoldSt1, Codec2} -> + fold_packets(Fun, FoldSt1, <<>>, Codec2); + {stop, FoldSt1, Codec2} -> + {ok, FoldSt1, Codec2} + end; + {incomplete, Codec1} -> + {ok, FoldSt0, Codec1} + end. + -spec is_empty(codec()) -> boolean(). is_empty(#codec{packet_buf = <<>>, crypto_buf = <<>>, tls_buf = <<>>}) -> true; is_empty(_) -> false. diff --git a/src/mtp_handler.erl b/src/mtp_handler.erl index 574ee43..a9a7524 100644 --- a/src/mtp_handler.erl +++ b/src/mtp_handler.erl @@ -275,21 +275,18 @@ handle_upstream_data(Bin, #state{stage = tunnel, {ok, S3#state{codec = UpCodec1}}; handle_upstream_data(Bin, #state{codec = Codec0} = S0) -> {ok, S, Codec} = - mtp_codec:fold_packets( + mtp_codec:fold_packets_if( fun(Decoded, S1, Codec1) -> case parse_upstream_data(Decoded, S1#state{codec = Codec1}) of {ok, S2} -> - {S2, S2#state.codec}; + {next, S2, S2#state.codec}; + {incomplete, S2} -> + {stop, S2, S2#state.codec}; {error, Err} -> error(Err) end end, S0, Bin, Codec0), - case mtp_codec:is_empty(Codec) of - true -> - {ok, S#state{codec = Codec}}; - false -> - handle_upstream_data(<<>>, S#state{codec = Codec}) - end. + {ok, S#state{codec = Codec}}. parse_upstream_data(<> = AllData, @@ -336,13 +333,13 @@ parse_upstream_data(<>, codec = Codec, stage = tunnel}, hibernate)); - {error, Reason} = Err -> + {error, Reason} when is_atom(Reason) -> mtp_metric:count_inc([?APP, protocol_error, total], 1, #{labels => [Reason]}), - Err + error({protocol_error, Reason, Header}) end; parse_upstream_data(Bin, #state{stage = Stage, codec = Codec0} = S) when Stage =/= tunnel -> Codec = mtp_codec:push_back(first, Bin, Codec0), - {ok, S#state{codec = Codec}}. + {incomplete, S#state{codec = Codec}}. assert_protocol(Protocol) -> {ok, AllowedProtocols} = application:get_env(?APP, allowed_protocols),