From fd57ec58721f02172a42b65dba47582b2035d4b3 Mon Sep 17 00:00:00 2001
From: Badlop <badlop@process-one.net>
Date: Mon, 31 Aug 2009 16:29:25 +0000
Subject: [PATCH] Support zlib compression after STARTTLS (thanks to Aleksey
 Shchepin)(EJAB-499)

SVN Revision: 2571
---
 doc/guide.html                      |  5 +----
 doc/guide.tex                       |  5 +----
 src/ejabberd_c2s.erl                | 15 ++++++++++++---
 src/ejabberd_zlib/ejabberd_zlib.erl | 19 ++++++++++++++++++-
 src/web/ejabberd_web_admin.erl      |  1 +
 5 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/doc/guide.html b/doc/guide.html
index a4ec1acb4..68876eace 100644
--- a/doc/guide.html
+++ b/doc/guide.html
@@ -779,10 +779,7 @@ password of one of the registered users who are granted access by the
 &#X2018;configure&#X2019; access rule.
 </DD><DT CLASS="dt-description"><B><TT>zlib</TT></B></DT><DD CLASS="dd-description"> This
 option specifies that Zlib stream compression (as defined in <A HREF="http://xmpp.org/extensions/xep-0138.html">XEP-0138</A>)
-is available on connections to the port. Client connections cannot use
-stream compression and stream encryption simultaneously. Hence, if you
-specify both <TT>starttls</TT> (or <TT>tls</TT>) and <TT>zlib</TT>, the latter
-option will not affect connections (there will be no stream compression).
+is available on connections to the port.
 </DD></DL><P>There are some additional global options that can be specified in the ejabberd configuration file (outside <TT>listen</TT>):
 </P><DL CLASS="description"><DT CLASS="dt-description">
 <B><TT>{s2s_use_starttls, true|false}</TT></B></DT><DD CLASS="dd-description">
diff --git a/doc/guide.tex b/doc/guide.tex
index b80b8be7d..32e3d67fe 100644
--- a/doc/guide.tex
+++ b/doc/guide.tex
@@ -920,10 +920,7 @@ This is a detailed description of each option allowed by the listening modules:
     `configure' access rule.
   \titem{zlib} \ind{options!zlib}\ind{protocols!XEP-0138: Stream Compression}\ind{Zlib}This
     option specifies that Zlib stream compression (as defined in \xepref{0138})
-    is available on connections to the port. Client connections cannot use
-    stream compression and stream encryption simultaneously. Hence, if you
-    specify both \option{starttls} (or \option{tls}) and \option{zlib}, the latter
-    option will not affect connections (there will be no stream compression).
+    is available on connections to the port.
 \end{description}
 
 There are some additional global options that can be specified in the ejabberd configuration file (outside \term{listen}):
diff --git a/src/ejabberd_c2s.erl b/src/ejabberd_c2s.erl
index 1c4a92bb8..88d26b1d5 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 3c3879b71..36beffedc 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 7b033315b..d0d27b8f8 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,