Support zlib compression after STARTTLS (thanks to Aleksey Shchepin)(EJAB-499)

SVN Revision: 2571

3 files changed, 31 insertions, 4 deletions

diff --git a/src/ejabberd_c2s.erl b/src/ejabberd_c2s.erl
index 1c4a92bb..88d26b1d 100644
--- a/src/ejabberd_c2s.erl
+++ b/src/ejabberd_c2s.erl
@@ -278,7 +278,8 @@ wait_for_stream({xmlstreamstart, _Name, Attrs}, StateData) ->
 				    Zlib = StateData#state.zlib,
 				    CompressFeature =
 					case Zlib andalso
-					    (SockMod == gen_tcp) of
+					    ((SockMod == gen_tcp) orelse
+					     (SockMod == tls)) of
 					    true ->
 						[{xmlelement, "compression",
 						  [{"xmlns", ?NS_FEATURE_COMPRESS}],
@@ -628,7 +629,8 @@ wait_for_feature_request({xmlstreamelement, El}, StateData) ->
 					   tls_enabled = true
 					  });
 	{?NS_COMPRESS, "compress"} when Zlib == true,
-					SockMod == gen_tcp ->
+					((SockMod == gen_tcp) or
+					 (SockMod == tls)) ->
 	    case xml:get_subtag(El, "method") of
 		false ->
 		    send_element(StateData,
@@ -1415,11 +1417,18 @@ get_auth_tags([_ | L], U, P, D, R) ->
 get_auth_tags([], U, P, D, R) ->
     {U, P, D, R}.
 
+%% Copied from ejabberd_socket.erl
+-record(socket_state, {sockmod, socket, receiver}).
+
 get_conn_type(StateData) ->
     case (StateData#state.sockmod):get_sockmod(StateData#state.socket) of
     gen_tcp -> c2s;
     tls -> c2s_tls;
-    ejabberd_zlib -> c2s_compressed;
+    ejabberd_zlib ->
+	case ejabberd_zlib:get_sockmod((StateData#state.socket)#socket_state.socket) of
+	    gen_tcp -> c2s_compressed;
+	    tls -> c2s_compressed_tls
+	end;
     ejabberd_http_poll -> http_poll;
     ejabberd_http_bind -> http_bind;
     _ -> unknown
diff --git a/src/ejabberd_zlib/ejabberd_zlib.erl b/src/ejabberd_zlib/ejabberd_zlib.erl
index 3c3879b7..36beffed 100644
--- a/src/ejabberd_zlib/ejabberd_zlib.erl
+++ b/src/ejabberd_zlib/ejabberd_zlib.erl
@@ -35,6 +35,7 @@
 	 recv/2, recv/3, recv_data/2,
 	 setopts/2,
 	 sockname/1, peername/1,
+	 get_sockmod/1,
 	 controlling_process/2,
 	 close/1]).
 
@@ -116,7 +117,20 @@ recv(#zlibsock{sockmod = SockMod, socket = Socket} = ZlibSock,
 	    Error
     end.
 
-recv_data(ZlibSock, Packet) ->
+recv_data(#zlibsock{sockmod = SockMod, socket = Socket} = ZlibSock, Packet) ->
+    case SockMod of
+	gen_tcp ->
+	    recv_data2(ZlibSock, Packet);
+	_ ->
+	    case SockMod:recv_data(Socket, Packet) of
+		{ok, Packet2} ->
+		    recv_data2(ZlibSock, Packet2);
+		Error ->
+		    Error
+	    end
+    end.
+
+recv_data2(ZlibSock, Packet) ->
     case catch recv_data1(ZlibSock, Packet) of
 	{'EXIT', Reason} ->
 	    {error, Reason};
@@ -158,6 +172,9 @@ sockname(#zlibsock{sockmod = SockMod, socket = Socket}) ->
 	    SockMod:sockname(Socket)
     end.
 
+get_sockmod(#zlibsock{sockmod = SockMod}) ->
+    SockMod.
+
 peername(#zlibsock{sockmod = SockMod, socket = Socket}) ->
     case SockMod of
 	gen_tcp ->
diff --git a/src/web/ejabberd_web_admin.erl b/src/web/ejabberd_web_admin.erl
index 7b033315..d0d27b8f 100644
--- a/src/web/ejabberd_web_admin.erl
+++ b/src/web/ejabberd_web_admin.erl
@@ -1663,6 +1663,7 @@ user_info(User, Server, Query, Lang) ->
 								 c2s -> "plain";
 								 c2s_tls -> "tls";
 								 c2s_compressed -> "zlib";
+								 c2s_compressed_tls -> "tls+zlib";
 								 http_bind -> "http-bind";
 								 http_poll -> "http-poll"
 							     end,