From cd88d342b99fdc710e3ef91313e54321b908237a Mon Sep 17 00:00:00 2001 From: Frank Diebolt Date: Wed, 30 Jan 2019 09:57:17 +0100 Subject: [PATCH] Apply shaping to websocket connections --- src/ejabberd_http_ws.erl | 10 +++--- src/ejabberd_websocket.erl | 62 +++++++++++++++++++++++++++++--------- 2 files changed, 53 insertions(+), 19 deletions(-) diff --git a/src/ejabberd_http_ws.erl b/src/ejabberd_http_ws.erl index 6e7a0901d..768a284ce 100644 --- a/src/ejabberd_http_ws.erl +++ b/src/ejabberd_http_ws.erl @@ -104,9 +104,8 @@ close({http_ws, FsmRef, _IP}) -> reset_stream({http_ws, _FsmRef, _IP} = Socket) -> Socket. -change_shaper({http_ws, _FsmRef, _IP}, _Shaper) -> - %% TODO??? - ok. +change_shaper({http_ws, FsmRef, _IP}, Shaper) -> + p1_fsm:send_all_state_event(FsmRef, {new_shaper, Shaper}). get_transport(_Socket) -> websocket. @@ -161,7 +160,10 @@ handle_event({activate, From}, StateName, State) -> end, Input), State#state{active = false, input = []} end, - {next_state, StateName, State1#state{c2s_pid = From}}. + {next_state, StateName, State1#state{c2s_pid = From}}; +handle_event({new_shaper, Shaper}, StateName, #state{ws = {_, WsPid}} = StateData) -> + WsPid ! {new_shaper, Shaper}, + {next_state, StateName, StateData}. handle_sync_event({send_xml, Packet}, _From, StateName, #state{ws = {_, WsPid}, rfc_compilant = R} = StateData) -> diff --git a/src/ejabberd_websocket.erl b/src/ejabberd_websocket.erl index 82443eec5..3c80cc82c 100644 --- a/src/ejabberd_websocket.erl +++ b/src/ejabberd_websocket.erl @@ -154,7 +154,7 @@ connect(#ws{socket = Socket, sockmod = SockMod} = Ws, WsLoop) -> _ -> SockMod:setopts(Socket, [{packet, 0}, {active, true}]) end, - ws_loop(none, Socket, WsHandleLoopPid, SockMod). + ws_loop(none, Socket, WsHandleLoopPid, SockMod, none). handshake(#ws{headers = Headers} = State) -> {_, Key} = lists:keyfind(<<"Sec-Websocket-Key">>, 1, @@ -187,18 +187,26 @@ find_subprotocol(Headers) -> end. -ws_loop(FrameInfo, Socket, WsHandleLoopPid, SocketMode) -> +ws_loop(FrameInfo, Socket, WsHandleLoopPid, SocketMode, Shaper) -> receive {DataType, _Socket, Data} when DataType =:= tcp orelse DataType =:= raw -> - case handle_data(DataType, FrameInfo, Data, Socket, WsHandleLoopPid, SocketMode) of + case handle_data(DataType, FrameInfo, Data, Socket, WsHandleLoopPid, SocketMode, Shaper) of {error, Error} -> ?DEBUG("TLS decode error ~p", [Error]), websocket_close(Socket, WsHandleLoopPid, SocketMode, 1002); % protocol error - {NewFrameInfo, ToSend} -> + {NewFrameInfo, ToSend, NewShaper} -> lists:foreach(fun(Pkt) -> SocketMode:send(Socket, Pkt) end, ToSend), - ws_loop(NewFrameInfo, Socket, WsHandleLoopPid, SocketMode) + ws_loop(NewFrameInfo, Socket, WsHandleLoopPid, SocketMode, NewShaper) end; + {new_shaper, NewShaper} -> + NewShaper = case NewShaper of + none when Shaper /= none -> + activate(Socket, SocketMode, true), none; + _ -> + NewShaper + end, + ws_loop(FrameInfo, Socket, WsHandleLoopPid, SocketMode, NewShaper); {tcp_closed, _Socket} -> ?DEBUG("TCP connection was closed, exit", []), websocket_close(Socket, WsHandleLoopPid, SocketMode, 0); @@ -220,15 +228,15 @@ ws_loop(FrameInfo, Socket, WsHandleLoopPid, SocketMode) -> {text, Data} -> SocketMode:send(Socket, encode_frame(Data, 1)), ws_loop(FrameInfo, Socket, WsHandleLoopPid, - SocketMode); + SocketMode, Shaper); {data, Data} -> SocketMode:send(Socket, encode_frame(Data, 2)), ws_loop(FrameInfo, Socket, WsHandleLoopPid, - SocketMode); + SocketMode, Shaper); {ping, Data} -> SocketMode:send(Socket, encode_frame(Data, 9)), ws_loop(FrameInfo, Socket, WsHandleLoopPid, - SocketMode); + SocketMode, Shaper); shutdown -> ?DEBUG("Shutdown request received, closing websocket " "with pid ~p", @@ -238,7 +246,7 @@ ws_loop(FrameInfo, Socket, WsHandleLoopPid, SocketMode) -> ?WARNING_MSG("Received unexpected message, ignoring: ~p", [_Ignored]), ws_loop(FrameInfo, Socket, WsHandleLoopPid, - SocketMode) + SocketMode, Shaper) end. encode_frame(Data, Opcode) -> @@ -393,17 +401,17 @@ process_frame(#frame_info{unprocessed = process_frame(FrameInfo#frame_info{unprocessed = <<>>}, <>). -handle_data(tcp, FrameInfo, Data, Socket, WsHandleLoopPid, fast_tls) -> +handle_data(tcp, FrameInfo, Data, Socket, WsHandleLoopPid, fast_tls, Shaper) -> case fast_tls:recv_data(Socket, Data) of {ok, NewData} -> - handle_data_int(FrameInfo, NewData, Socket, WsHandleLoopPid, fast_tls); + handle_data_int(FrameInfo, NewData, Socket, WsHandleLoopPid, fast_tls, Shaper); {error, Error} -> {error, Error} end; -handle_data(_, FrameInfo, Data, Socket, WsHandleLoopPid, SockMod) -> - handle_data_int(FrameInfo, Data, Socket, WsHandleLoopPid, SockMod). +handle_data(_, FrameInfo, Data, Socket, WsHandleLoopPid, SockMod, Shaper) -> + handle_data_int(FrameInfo, Data, Socket, WsHandleLoopPid, SockMod, Shaper). -handle_data_int(FrameInfo, Data, _Socket, WsHandleLoopPid, _SocketMode) -> +handle_data_int(FrameInfo, Data, Socket, WsHandleLoopPid, SocketMode, Shaper) -> {NewFrameInfo, Recv, Send} = process_frame(FrameInfo, Data), lists:foreach(fun (El) -> case El of @@ -416,7 +424,7 @@ handle_data_int(FrameInfo, Data, _Socket, WsHandleLoopPid, _SocketMode) -> end end, Recv), - {NewFrameInfo, Send}. + {NewFrameInfo, Send, handle_shaping(Data, Socket, SocketMode, Shaper)}. websocket_close(Socket, WsHandleLoopPid, SocketMode, CloseCode) when CloseCode > 0 -> @@ -429,3 +437,27 @@ websocket_close(Socket, WsHandleLoopPid, SocketMode, _CloseCode) -> get_origin() -> ejabberd_option:websocket_origin(). + +handle_shaping(_Data, _Socket, _SocketMode, none) -> + none; +handle_shaping(Data, Socket, SocketMode, Shaper) -> + {NewShaper, Pause} = ejabberd_shaper:update(Shaper, byte_size(Data)), + if Pause > 0 -> + activate_after(Socket, self(), Pause); + true -> activate(Socket, SocketMode, once) + end, + NewShaper. + +activate(Socket, SockMod, ActiveState) -> + case SockMod of + gen_tcp -> inet:setopts(Socket, [{active, ActiveState}]); + _ -> SockMod:setopts(Socket, [{active, ActiveState}]) + end. + +activate_after(Socket, Pid, Pause) -> + if Pause > 0 -> + erlang:send_after(Pause, Pid, {tcp, Socket, <<>>}); + true -> + Pid ! {tcp, Socket, <<>>} + end, + ok.