Merge pull request #167 from esl/zlib_cwe400_fix

Add zlib inflated data size limit

Author: michalwski

apps/ejabberd/c_src/ejabberd_zlib_drv.c

@@ -109,79 +109,86 @@ static void ejabberd_zlib_drv_stop(ErlDrvData handle)
 
 
 static ErlDrvSSizeT ejabberd_zlib_drv_control(ErlDrvData handle,
-				     unsigned int command,
-				     char *buf, ErlDrvSizeT len,
-				     char **rbuf, ErlDrvSizeT rlen)
+        unsigned int command,
+        char *buf, ErlDrvSizeT len,
+        char **rbuf, ErlDrvSizeT rlen)
 {
-   ejabberd_zlib_data *d = (ejabberd_zlib_data *)handle;
-   int err;
-   int size;
-   ErlDrvBinary *b;
-
-   switch (command)
-   {
-      case DEFLATE:
-	 size = BUF_SIZE + 1;
-	 rlen = 1;
-	 b = driver_alloc_binary(size);
-	 b->orig_bytes[0] = 0;
-
-	 d->d_stream->next_in = (unsigned char *)buf;
-	 d->d_stream->avail_in = len;
-	 d->d_stream->avail_out = 0;
-	 err = Z_OK;
-
-	 while (err == Z_OK && d->d_stream->avail_out == 0)
-	 {
-	    d->d_stream->next_out = (unsigned char *)b->orig_bytes + rlen;
-	    d->d_stream->avail_out = BUF_SIZE;
-
-	    err = deflate(d->d_stream, Z_SYNC_FLUSH);
-	    die_unless((err == Z_OK) || (err == Z_STREAM_END),
-		       "Deflate error");
-
-	    rlen += (BUF_SIZE - d->d_stream->avail_out);
-	    size += (BUF_SIZE - d->d_stream->avail_out);
-	    b = driver_realloc_binary(b, size);
-	 }
-	 b = driver_realloc_binary(b, rlen);
-	 *rbuf = (char *)b;
-	 return rlen;
-      case INFLATE:
-	 size = BUF_SIZE + 1;
-	 rlen = 1;
-	 b = driver_alloc_binary(size);
-	 b->orig_bytes[0] = 0;
-
-	 if (len > 0) {
-	    d->i_stream->next_in = (unsigned char *)buf;
-	    d->i_stream->avail_in = len;
-	    d->i_stream->avail_out = 0;
-	    err = Z_OK;
-
-	    while (err == Z_OK && d->i_stream->avail_out == 0)
-	    {
-	       d->i_stream->next_out = (unsigned char *)b->orig_bytes + rlen;
-	       d->i_stream->avail_out = BUF_SIZE;
-
-	       err = inflate(d->i_stream, Z_SYNC_FLUSH);
-	       die_unless((err == Z_OK) || (err == Z_STREAM_END),
-			  "Inflate error");
-
-	       rlen += (BUF_SIZE - d->i_stream->avail_out);
-	       size += (BUF_SIZE - d->i_stream->avail_out);
-	       b = driver_realloc_binary(b, size);
-	    }
-	 }
-	 b = driver_realloc_binary(b, rlen);
-	 *rbuf = (char *)b;
-	 return rlen;
-   }
-
-   b = driver_alloc_binary(1);
-   b->orig_bytes[0] = 0;
-   *rbuf = (char *)b;
-   return 1;
+    ejabberd_zlib_data *d = (ejabberd_zlib_data *)handle;
+    int err;
+    int size;
+    int size_limit;
+    ErlDrvBinary *b;
+
+    // operation is in command's 2 lower bits and size_limit is in bits higher than 1
+    size_limit = command >> 2; // applies only to inflation
+    command = command & 3;
+    switch (command)
+    {
+        case DEFLATE:
+            size = BUF_SIZE + 1;
+            rlen = 1;
+            b = driver_alloc_binary(size);
+            b->orig_bytes[0] = 0;
+
+            d->d_stream->next_in = (unsigned char *)buf;
+            d->d_stream->avail_in = len;
+            d->d_stream->avail_out = 0;
+            err = Z_OK;
+
+            while (err == Z_OK && d->d_stream->avail_out == 0)
+            {
+                d->d_stream->next_out = (unsigned char *)b->orig_bytes + rlen;
+                d->d_stream->avail_out = BUF_SIZE;
+
+                err = deflate(d->d_stream, Z_SYNC_FLUSH);
+                die_unless((err == Z_OK) || (err == Z_STREAM_END),
+                        "deflate_error");
+
+                rlen += (BUF_SIZE - d->d_stream->avail_out);
+                size += (BUF_SIZE - d->d_stream->avail_out);
+                b = driver_realloc_binary(b, size);
+            }
+            b = driver_realloc_binary(b, rlen);
+            *rbuf = (char *)b;
+            return rlen;
+        case INFLATE:
+            size = BUF_SIZE + 1;
+            rlen = 1;
+            b = driver_alloc_binary(size);
+            b->orig_bytes[0] = 0;
+
+            if (len > 0) {
+                d->i_stream->next_in = (unsigned char *)buf;
+                d->i_stream->avail_in = len;
+                d->i_stream->avail_out = 0;
+                err = Z_OK;
+
+                while (err == Z_OK && d->i_stream->avail_out == 0)
+                {
+                    d->i_stream->next_out = (unsigned char *)b->orig_bytes + rlen;
+                    d->i_stream->avail_out = BUF_SIZE;
+
+                    err = inflate(d->i_stream, Z_SYNC_FLUSH);
+                    die_unless((err == Z_OK) || (err == Z_STREAM_END),
+                            "inflate_error");
+
+                    rlen += (BUF_SIZE - d->i_stream->avail_out);
+                    die_unless((rlen < size_limit) || (size_limit == 0),
+                            "inflate_size_exceeded");
+
+                    size += (BUF_SIZE - d->i_stream->avail_out);
+                    b = driver_realloc_binary(b, size);
+                }
+            }
+            b = driver_realloc_binary(b, rlen);
+            *rbuf = (char *)b;
+            return rlen;
+    }
+
+    b = driver_alloc_binary(1);
+    b->orig_bytes[0] = 0;
+    *rbuf = (char *)b;
+    return 1;
 }
 
 

apps/ejabberd/src/ejabberd_c2s.erl

@@ -80,7 +80,7 @@
 		sasl_state,
 		access,
 		shaper,
-		zlib = false,
+		zlib = {false, 0},
 		tls = false,
 		tls_required = false,
 		tls_enabled = false,
@@ -161,8 +161,8 @@ get_presence(FsmRef) ->
     ?GEN_FSM:sync_send_all_state_event(FsmRef, get_presence, 1000).
 
 get_aux_field(Key, #state{aux_fields = Opts}) ->
-    case lists:keysearch(Key, 1, Opts) of
-	{value, {_, Val}} ->
+    case lists:keyfind(Key, 1, Opts) of
+	{_, Val} ->
 	    {ok, Val};
 	_ ->
 	    error
@@ -208,20 +208,23 @@ stop(FsmRef) ->
 %%          {stop, StopReason}
 %%----------------------------------------------------------------------
 init([{SockMod, Socket}, Opts]) ->
-    Access = case lists:keysearch(access, 1, Opts) of
-		 {value, {_, A}} -> A;
+    Access = case lists:keyfind(access, 1, Opts) of
+		 {_, A} -> A;
 		 _ -> all
 	     end,
-    Shaper = case lists:keysearch(shaper, 1, Opts) of
-		 {value, {_, S}} -> S;
+    Shaper = case lists:keyfind(shaper, 1, Opts) of
+		 {_, S} -> S;
 		 _ -> none
 	     end,
     XMLSocket =
-	case lists:keysearch(xml_socket, 1, Opts) of
-	    {value, {_, XS}} -> XS;
+	case lists:keyfind(xml_socket, 1, Opts) of
+	    {_, XS} -> XS;
 	    _ -> false
 	end,
-    Zlib = lists:member(zlib, Opts),
+    Zlib = case lists:keyfind(zlib, 1, Opts) of
+               {_, ZlibLimit} -> {true, ZlibLimit};
+               _ -> {false, 0}
+           end,
     StartTLS = lists:member(starttls, Opts),
     StartTLSRequired = lists:member(starttls_required, Opts),
     TLSEnabled = lists:member(tls, Opts),
@@ -328,7 +331,7 @@ wait_for_stream({xmlstreamstart, _Name, Attrs}, StateData) ->
 		    SockMod =
 			 (StateData#state.sockmod):get_sockmod(
 			   StateData#state.socket),
-		    Zlib = StateData#state.zlib,
+		    {Zlib, _} = StateData#state.zlib,
 		    CompressFeature =
 			case Zlib andalso
 			      ((SockMod == gen_tcp) orelse
@@ -469,8 +472,7 @@ wait_for_stream(closed, StateData) ->
 wait_for_auth({xmlstreamelement, El}, StateData) ->
     case is_auth_packet(El) of
 	{auth, _ID, get, {U, _, _, _}} ->
-	    XE = #xmlel{name = Name,
-                        attrs = Attrs} = jlib:make_result_iq_reply(El),
+	    XE = jlib:make_result_iq_reply(El),
 	    case U of
 		<<>> ->
 		    UCdata = [];
@@ -613,7 +615,7 @@ wait_for_auth(closed, StateData) ->
 
 wait_for_feature_request({xmlstreamelement, El}, StateData) ->
     #xmlel{name = Name, attrs = Attrs, children = Els} = El,
-    Zlib = StateData#state.zlib,
+    {Zlib, ZlibLimit} = StateData#state.zlib,
     TLS = StateData#state.tls,
     TLSEnabled = StateData#state.tls_enabled,
     TLSRequired = StateData#state.tls_required,
@@ -705,7 +707,7 @@ wait_for_feature_request({xmlstreamelement, El}, StateData) ->
 			<<"zlib">> ->
 			    Socket = StateData#state.socket,
 			    ZlibSocket = (StateData#state.sockmod):compress(
-					   Socket,
+					   Socket, ZlibLimit,
 					   xml:element_to_binary(
 					      #xmlel{name = <<"compressed">>,
 					             attrs = [{<<"xmlns">>, ?NS_COMPRESS}]})),
@@ -979,7 +981,7 @@ session_established({xmlstreamend, _Name}, StateData) ->
     send_trailer(StateData),
     {stop, normal, StateData};
 
-session_established({xmlstreamerror, "XML stanza is too big" = E}, StateData) ->
+session_established({xmlstreamerror, <<"XML stanza is too big">> = E}, StateData) ->
     send_element(StateData, ?POLICY_VIOLATION_ERR(StateData#state.lang, E)),
     send_trailer(StateData),
     {stop, normal, StateData};
@@ -2217,8 +2219,8 @@ check_from(El, FromJID) ->
     end.
 
 fsm_limit_opts(Opts) ->
-    case lists:keysearch(max_fsm_queue, 1, Opts) of
-	{value, {_, N}} when is_integer(N) ->
+    case lists:keyfind(max_fsm_queue, 1, Opts) of
+	{_, N} when is_integer(N) ->
 	    [{max_queue, N}];
 	_ ->
 	    case ejabberd_config:get_local_option(max_fsm_queue) of

apps/ejabberd/src/ejabberd_receiver.erl

@@ -56,6 +56,7 @@
 		timeout}).
 
 -define(HIBERNATE_TIMEOUT, 90000).
+-define(GEN_FSM, p1_fsm).
 
 %%====================================================================
 %% API
@@ -160,7 +161,10 @@ handle_call({compress, ZlibSocket}, _From,
     case ejabberd_zlib:recv_data(ZlibSocket, "") of
 	{ok, ZlibData} ->
 	    {reply, ok, process_data(ZlibData, NewState), ?HIBERNATE_TIMEOUT};
-	{error, _Reason} ->
+	{error, inflate_size_exceeded} ->
+	    ?GEN_FSM:send_event(C2SPid, {xmlstreamerror, <<"XML stanza is too big">>}),
+	    {reply, ok, NewState, ?HIBERNATE_TIMEOUT};
+	{error, inflate_error} ->
 	    {stop, normal, ok, NewState}
     end;
 handle_call(reset_stream, _From,
@@ -205,6 +209,7 @@ handle_cast(_Msg, State) ->
 %%--------------------------------------------------------------------
 handle_info({Tag, _TCPSocket, Data},
 	    #state{socket = Socket,
+		   c2s_pid = C2SPid,
 		   sock_mod = SockMod} = State)
   when (Tag == tcp) or (Tag == ssl) or (Tag == ejabberd_xml) ->
     case SockMod of
@@ -221,7 +226,10 @@ handle_info({Tag, _TCPSocket, Data},
 		{ok, ZlibData} ->
 		    {noreply, process_data(ZlibData, State),
 		     ?HIBERNATE_TIMEOUT};
-		{error, _Reason} ->
+		{error, inflate_size_exceeded} ->
+		    ?GEN_FSM:send_event(C2SPid, {xmlstreamerror, <<"XML stanza is too big">>}),
+		    {noreply, State, ?HIBERNATE_TIMEOUT};
+		{error, inflate_error} ->
 		    {stop, normal, State}
 	    end;
 	_ ->

apps/ejabberd/src/ejabberd_socket.erl

@@ -34,7 +34,7 @@
 	 starttls/2,
 	 starttls/3,
 	 compress/1,
-	 compress/2,
+	 compress/3,
 	 reset_stream/1,
 	 send/2,
 	 send_xml/2,
@@ -152,10 +152,11 @@ compress(SocketData) ->
     ejabberd_receiver:compress(SocketData#socket_state.receiver, ZlibSocket),
     SocketData#socket_state{socket = ZlibSocket, sockmod = ejabberd_zlib}.
 
-compress(SocketData, Data) ->
+compress(SocketData, InflateSizeLimit, Data) ->
     {ok, ZlibSocket} = ejabberd_zlib:enable_zlib(
 			 SocketData#socket_state.sockmod,
-			 SocketData#socket_state.socket),
+			 SocketData#socket_state.socket,
+			 InflateSizeLimit),
     ejabberd_receiver:compress(SocketData#socket_state.receiver, ZlibSocket),
     send(SocketData, Data),
     SocketData#socket_state{socket = ZlibSocket, sockmod = ejabberd_zlib}.

apps/ejabberd/src/ejabberd_zlib.erl

@@ -30,7 +30,7 @@
 -behaviour(gen_server).
 
 -export([start/0, start_link/0,
-	 enable_zlib/2, disable_zlib/1,
+	 enable_zlib/3, disable_zlib/1,
 	 send/2,
 	 recv/2, recv/3, recv_data/2,
 	 setopts/2,
@@ -50,7 +50,7 @@
 -define(DEFLATE, 1).
 -define(INFLATE, 2).
 
--record(zlibsock, {sockmod, socket, zlibport}).
+-record(zlibsock, {sockmod, socket, zlibport, inflate_size_limit = 0}).
 
 start() ->
     gen_server:start({local, ?MODULE}, ?MODULE, [], []).
@@ -94,13 +94,14 @@ terminate(_Reason, Port) ->
     ok.
 
 
-enable_zlib(SockMod, Socket) ->
+enable_zlib(SockMod, Socket, InflateSizeLimit) ->
     case erl_ddll:load_driver(ejabberd:get_so_path(), ejabberd_zlib_drv) of
 	ok -> ok;
 	{error, already_loaded} -> ok
     end,
     Port = open_port({spawn, ejabberd_zlib_drv}, [binary]),
-    {ok, #zlibsock{sockmod = SockMod, socket = Socket, zlibport = Port}}.
+    {ok, #zlibsock{sockmod = SockMod, socket = Socket, zlibport = Port,
+                   inflate_size_limit = InflateSizeLimit}}.
 
 disable_zlib(#zlibsock{sockmod = SockMod, socket = Socket, zlibport = Port}) ->
     port_close(Port),
@@ -138,12 +139,12 @@ recv_data2(ZlibSock, Packet) ->
 	    Res
     end.
 
-recv_data1(#zlibsock{zlibport = Port} = _ZlibSock, Packet) ->
-    case port_control(Port, ?INFLATE, Packet) of
+recv_data1(#zlibsock{zlibport = Port, inflate_size_limit = SizeLimit} = _ZlibSock, Packet) ->
+    case port_control(Port, SizeLimit bsl 2 + ?INFLATE, Packet) of
 	<<0, In/binary>> ->
 	    {ok, In};
 	<<1, Error/binary>> ->
-	    {error, binary_to_list(Error)}
+	    {error, erlang:binary_to_existing_atom(Error, utf8)}
     end.
 
 send(#zlibsock{sockmod = SockMod, socket = Socket, zlibport = Port},
@@ -152,7 +153,9 @@ send(#zlibsock{sockmod = SockMod, socket = Socket, zlibport = Port},
 	<<0, Out/binary>> ->
 	    SockMod:send(Socket, Out);
 	<<1, Error/binary>> ->
-	    {error, binary_to_list(Error)}
+	    {error, erlang:binary_to_existing_atom(Error, utf8)};
+	_ ->
+	    {error, deflate_error}
     end.