plugins/mod_compression.lua

   1 -- Prosody IM
   2 -- Copyright (C) 2009 Tobias Markmann
   3 --
   4 -- This project is MIT/X11 licensed. Please see the
   5 -- COPYING file in the source package for more information.
   6 --
   7
   8 local st = require "util.stanza";
   9 local zlib = require "zlib";
  10 local pcall = pcall;
  11 local xmlns_compression_feature = "http://jabber.org/features/compress"
  12 local xmlns_compression_protocol = "http://jabber.org/protocol/compress"
  13 local xmlns_stream = "http://etherx.jabber.org/streams";
  14 local compression_stream_feature = st.stanza("compression", {xmlns=xmlns_compression_feature}):tag("method"):text("zlib"):up();
  15
  16 local compression_level = module:get_option("compression_level");
  17 -- if not defined assume admin wants best compression
  18 if compression_level == nil then compression_level = 9 end;
  19
  20
  21 compression_level = tonumber(compression_level);
  22 if not compression_level or compression_level < 1 or compression_level > 9 then
  23         module:log("warn", "Invalid compression level in config: %s", tostring(compression_level));
  24         module:log("warn", "Module loading aborted. Compression won't be available.");
  25         return;
  26 end
  27
  28 module:add_event_hook("stream-features",
  29                 function (session, features)
  30                         if not session.compressed then
  31                                 -- FIXME only advertise compression support when TLS layer has no compression enabled
  32                                 features:add_child(compression_stream_feature);
  33                         end
  34                 end
  35 );
  36
  37 module:hook("s2s-stream-features",
  38                 function (data)
  39                         local session, features = data.session, data.features;
  40                         -- FIXME only advertise compression support when TLS layer has no compression enabled
  41                         if not session.compressed then
  42                                 features:add_child(compression_stream_feature);
  43                         end
  44                 end
  45 );
  46
  47 -- Hook to activate compression if remote server supports it.
  48 module:hook_stanza(xmlns_stream, "features",
  49                 function (session, stanza)
  50                         if not session.compressed then
  51                                 -- does remote server support compression?
  52                                 local comp_st = stanza:child_with_name("compression");
  53                                 if comp_st then
  54                                         -- do we support the mechanism
  55                                         for a in comp_st:children() do
  56                                                 local algorithm = a[1]
  57                                                 if algorithm == "zlib" then
  58                                                         session.sends2s(st.stanza("compress", {xmlns=xmlns_compression_protocol}):tag("method"):text("zlib"))
  59                                                         session.log("info", "Enabled compression using zlib.")
  60                                                         return true;
  61                                                 end
  62                                         end
  63                                         session.log("debug", "Remote server supports no compression algorithm we support.")
  64                                 end
  65                         end
  66                 end
  67 , 250);
  68
  69
  70 -- returns either nil or a fully functional ready to use inflate stream
  71 local function get_deflate_stream(session)
  72         local status, deflate_stream = pcall(zlib.deflate, compression_level);
  73         if status == false then
  74                 local error_st = st.stanza("failure", {xmlns=xmlns_compression_protocol}):tag("setup-failed");
  75                 (session.sends2s or session.send)(error_st);
  76                 session.log("error", "Failed to create zlib.deflate filter.");
  77                 module:log("error", deflate_stream);
  78                 return
  79         end
  80         return deflate_stream
  81 end
  82
  83 -- returns either nil or a fully functional ready to use inflate stream
  84 local function get_inflate_stream(session)
  85         local status, inflate_stream = pcall(zlib.inflate);
  86         if status == false then
  87                 local error_st = st.stanza("failure", {xmlns=xmlns_compression_protocol}):tag("setup-failed");
  88                 (session.sends2s or session.send)(error_st);
  89                 session.log("error", "Failed to create zlib.deflate filter.");
  90                 module:log("error", inflate_stream);
  91                 return
  92         end
  93         return inflate_stream
  94 end
  95
  96 -- setup compression for a stream
  97 local function setup_compression(session, deflate_stream)
  98         local old_send = (session.sends2s or session.send);
  99
 100         local new_send = function(t)
 101                         --TODO: Better code injection in the sending process
 102                         session.log(t)
 103                         local status, compressed, eof = pcall(deflate_stream, tostring(t), 'sync');
 104                         if status == false then
 105                                 session:close({
 106                                         condition = "undefined-condition";
 107                                         text = compressed;
 108                                         extra = st.stanza("failure", {xmlns="http://jabber.org/protocol/compress"}):tag("processing-failed");
 109                                 });
 110                                 module:log("warn", compressed);
 111                                 return;
 112                         end
 113                         session.conn:write(compressed);
 114                 end;
 115
 116         if session.sends2s then session.sends2s = new_send
 117         elseif session.send then session.send = new_send end
 118 end
 119
 120 -- setup decompression for a stream
 121 local function setup_decompression(session, inflate_stream)
 122         local old_data = session.data
 123         session.data = function(conn, data)
 124                         local status, decompressed, eof = pcall(inflate_stream, data);
 125                         if status == false then
 126                                 session:close({
 127                                         condition = "undefined-condition";
 128                                         text = decompressed;
 129                                         extra = st.stanza("failure", {xmlns="http://jabber.org/protocol/compress"}):tag("processing-failed");
 130                                 });
 131                                 module:log("warn", decompressed);
 132                                 return;
 133                         end
 134                         old_data(conn, decompressed);
 135                 end;
 136 end
 137
 138 module:add_handler({"s2sout_unauthed", "s2sout"}, "compressed", xmlns_compression_protocol,
 139                 function(session ,stanza)
 140                         session.log("debug", "Activating compression...")
 141                         -- create deflate and inflate streams
 142                         local deflate_stream = get_deflate_stream(session);
 143                         if not deflate_stream then return end
 144
 145                         local inflate_stream = get_inflate_stream(session);
 146                         if not inflate_stream then return end
 147
 148                         -- setup compression for session.w
 149                         setup_compression(session, deflate_stream);
 150
 151                         -- setup decompression for session.data
 152                         setup_decompression(session, inflate_stream);
 153                         local session_reset_stream = session.reset_stream;
 154                         session.reset_stream = function(session)
 155                                         session_reset_stream(session);
 156                                         setup_decompression(session, inflate_stream);
 157                                         return true;
 158                                 end;
 159                         session:reset_stream();
 160                         local default_stream_attr = {xmlns = "jabber:server", ["xmlns:stream"] = "http://etherx.jabber.org/streams",
 161                                                                                 ["xmlns:db"] = 'jabber:server:dialback', version = "1.0", to = session.to_host, from = session.from_host};
 162                         session.sends2s("<?xml version='1.0'?>");
 163                         session.sends2s(st.stanza("stream:stream", default_stream_attr):top_tag());
 164                         session.compressed = true;
 165                 end
 166 );
 167
 168 module:add_handler({"c2s_unauthed", "c2s", "s2sin_unauthed", "s2sin"}, "compress", xmlns_compression_protocol,
 169                 function(session, stanza)
 170                         -- fail if we are already compressed
 171                         if session.compressed then
 172                                 local error_st = st.stanza("failure", {xmlns=xmlns_compression_protocol}):tag("unsupported-method");
 173                                 (session.sends2s or session.send)(error_st);
 174                                 session.log("warn", "Tried to establish another compression layer.");
 175                         end
 176
 177                         -- checking if the compression method is supported
 178                         local method = stanza:child_with_name("method")[1];
 179                         if method == "zlib" then
 180                                 session.log("debug", method.." compression selected.");
 181
 182                                 -- create deflate and inflate streams
 183                                 local deflate_stream = get_deflate_stream(session);
 184                                 if not deflate_stream then return end
 185
 186                                 local inflate_stream = get_inflate_stream(session);
 187                                 if not inflate_stream then return end
 188
 189                                 (session.sends2s or session.send)(st.stanza("compressed", {xmlns=xmlns_compression_protocol}));
 190                                 session:reset_stream();
 191
 192                                 -- setup compression for session.w
 193                                 setup_compression(session, deflate_stream);
 194
 195                                 -- setup decompression for session.data
 196                                 setup_decompression(session, inflate_stream);
 197
 198                                 local session_reset_stream = session.reset_stream;
 199                                 session.reset_stream = function(session)
 200                                                 session_reset_stream(session);
 201                                                 setup_decompression(session, inflate_stream);
 202                                                 return true;
 203                                         end;
 204                                 session.compressed = true;
 205                         else
 206                                 session.log("warn", method.." compression selected. But we don't support it.");
 207                                 local error_st = st.stanza("failure", {xmlns=xmlns_compression_protocol}):tag("unsupported-method");
 208                                 (session.sends2s or session.send)(error_st);
 209                         end
 210                 end
 211 );
 212