Apply shaping to websocket connections

2 files changed, 52 insertions, 18 deletions

diff --git a/src/ejabberd_http_ws.erl b/src/ejabberd_http_ws.erl
index 4b54e67ec..79406c672 100644
--- a/src/ejabberd_http_ws.erl
+++ b/src/ejabberd_http_ws.erl
@@ -108,9 +108,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.
@@ -169,7 +168,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 506ff142b..04e7f916a 100644
--- a/src/ejabberd_websocket.erl
+++ b/src/ejabberd_websocket.erl
@@ -141,7 +141,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,
@@ -174,18 +174,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);
@@ -204,11 +212,11 @@ ws_loop(FrameInfo, Socket, WsHandleLoopPid, SocketMode) ->
         {send, Data} ->
             SocketMode:send(Socket, encode_frame(Data, 1)),
             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",
@@ -218,7 +226,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) ->
@@ -373,17 +381,17 @@ process_frame(#frame_info{unprocessed =
     process_frame(FrameInfo#frame_info{unprocessed = <<>>},
                   <<UnprocessedPre/binary, Data/binary>>).
 
-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
@@ -396,7 +404,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 ->
@@ -406,3 +414,27 @@ websocket_close(Socket, WsHandleLoopPid,
 websocket_close(Socket, WsHandleLoopPid, SocketMode, _CloseCode) ->
     WsHandleLoopPid ! closed,
     SocketMode:close(Socket).
+
+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.