add handle_info to optional callback functions

* Module:handle_info(Info, CBState) -> {noreply, NewCbState} | {noreply, NewCbState, Timeout} | {close, Reason, NewCbState}
erlyaws · Dec 11, 2012 · 167fdb8 · 167fdb8
1 parent e755418
commit 167fdb8
Showing 1 changed file with 28 additions and 5 deletions.
diff --git a/src/yaws_websockets.erl b/src/yaws_websockets.erl
@@ -41,7 +41,8 @@
                  terminate_fun,
                  open_fun,
                  msg_fun_1,
-                 msg_fun_2}).
+                 msg_fun_2,
+		 info_fun}).
 
 -export([start/3, send/2]).
 -export([init/1, handle_call/3, handle_cast/2, handle_info/2,
@@ -277,9 +278,26 @@ handle_info(timeout, #state{wsstate=WSState}=State) ->
             {stop, normal, State#state{reason={error,Reason}}}
     end;
 
-%% Skip all other messages
-handle_info(_Info, State) ->
-    {noreply, State, State#state.timeout}.
+%% handle info messages
+handle_info(Info, State) ->
+    case (CbInfo=State#state.cbinfo)#cbinfo.info_fun of
+	undefined ->
+	    {noreply, State, State#state.timeout};
+	InfoFun ->
+	    CbMod = CbInfo#cbinfo.module,
+	    {FrameState, CbState} = State#state.cbstate,
+	    Res = CbMod:InfoFun(Info, CbState),
+		case Res of
+		    {noreply, NewCbState} ->
+			{noreply, State#state{cbstate={FrameState,NewCbState}}};
+		    {noreply, NewCbState, Timeout} ->
+			{noreply, State#state{cbstate={FrameState,NewCbState}}, Timeout};
+		    {close, Reason, NewCbState} ->
+			{stop, normal, State#state{reason=Reason,
+                                       cbstate={FrameState,NewCbState}}}
+		    end
+    end.
+
 
 
 %% ----
@@ -334,12 +352,17 @@ get_callback_info(Mod) ->
                           true  -> handle_message;
                           false -> undefined
                       end,
+	    InfoFun = case erlang:function_exported(Mod, handle_info, 2) of
+                          true  -> handle_info;
+                          false -> undefined
+                      end,
             {ok, #cbinfo{module        = Mod,
                          init_fun      = InitFun,
                          terminate_fun = TerminateFun,
                          open_fun      = OpenFun,
                          msg_fun_1     = MsgFun1,
-                         msg_fun_2     = MsgFun2}};
+                         msg_fun_2     = MsgFun2,
+			 info_fun      = InfoFun}};
         {error, Reason} ->
             error_logger:error_msg("Cannot load callback module '~p': ~p",
                                    [Mod, Reason]),