src/modules/extra/m_ziplink.cpp

   1 #include <string>
   2 #include <vector>
   3
   4 #include "zlib.h"
   5
   6 #include "inspircd_config.h"
   7 #include "configreader.h"
   8 #include "users.h"
   9 #include "channels.h"
  10 #include "modules.h"
  11
  12 #include "socket.h"
  13 #include "hashcomp.h"
  14 #include "inspircd.h"
  15
  16 #include "transport.h"
  17
  18 /* $ModDesc: Provides zlib link support for servers */
  19 /* $LinkerFlags: -lz */
  20 /* $ModDep: transport.h */
  21
  22 /*
  23  * Compressed data is transmitted across the link in the following format:
  24  *
  25  *   0   1   2   3   4 ... n
  26  * +---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+
  27  * |       n       |              Z0 -> Zn                         |
  28  * +---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+
  29  *
  30  * Where: n is the size of a frame, in network byte order, 4 bytes.
  31  * Z0 through Zn are Zlib compressed data, n bytes in length.
  32  *
  33  * If the module fails to read the entire frame, then it will buffer
  34  * the portion of the last frame it received, then attempt to read
  35  * the next part of the frame next time a write notification arrives.
  36  *
  37  * ZLIB_BEST_COMPRESSION (9) is used for all sending of data with
  38  * a flush after each frame. A frame may contain multiple lines
  39  * and should be treated as raw binary data.
  40  *
  41  */
  42
  43
  44 enum izip_status { IZIP_WAITFIRST, IZIP_OPEN, IZIP_CLOSED };
  45
  46 const unsigned int CHUNK = 16384;
  47
  48 /** Represents an ZIP user's extra data
  49  */
  50 class izip_session : public classbase
  51 {
  52  public:
  53         z_stream c_stream; /* compression stream */
  54         z_stream d_stream; /* decompress stream */
  55         izip_status status;
  56         int need_bytes;
  57         int fd;
  58         std::string inbuf;
  59 };
  60
  61 class ModuleZLib : public Module
  62 {
  63         izip_session sessions[MAX_DESCRIPTORS];
  64         float total_out_compressed;
  65         float total_in_compressed;
  66         float total_out_uncompressed;
  67         float total_in_uncompressed;
  68
  69  public:
  70
  71         ModuleZLib(InspIRCd* Me)
  72                 : Module::Module(Me)
  73         {
  74                 ServerInstance->PublishInterface("InspSocketHook", this);
  75
  76                 total_out_compressed = total_in_compressed = 0;
  77                 total_out_uncompressed = total_out_uncompressed = 0;
  78         }
  79
  80         virtual ~ModuleZLib()
  81         {
  82         }
  83
  84         virtual Version GetVersion()
  85         {
  86                 return Version(1, 1, 0, 0, VF_VENDOR, API_VERSION);
  87         }
  88
  89         void Implements(char* List)
  90         {
  91                 List[I_OnRawSocketConnect] = List[I_OnRawSocketAccept] = List[I_OnRawSocketClose] = List[I_OnRawSocketRead] = List[I_OnRawSocketWrite] = 1;
  92                 List[I_OnRequest] = 1;
  93         }
  94
  95         virtual char* OnRequest(Request* request)
  96         {
  97                 ISHRequest* ISR = (ISHRequest*)request;
  98                 if (strcmp("IS_NAME", request->GetId()) == 0)
  99                 {
 100                         return "zip";
 101                 }
 102                 else if (strcmp("IS_HOOK", request->GetId()) == 0)
 103                 {
 104                         return ServerInstance->Config->AddIOHook((Module*)this, (InspSocket*)ISR->Sock) ? (char*)"OK" : NULL;
 105                 }
 106                 else if (strcmp("IS_UNHOOK", request->GetId()) == 0)
 107                 {
 108                         return ServerInstance->Config->DelIOHook((InspSocket*)ISR->Sock) ? (char*)"OK" : NULL;
 109                 }
 110                 else if (strcmp("IS_HSDONE", request->GetId()) == 0)
 111                 {
 112                         return "OK";
 113                 }
 114                 else if (strcmp("IS_ATTACH", request->GetId()) == 0)
 115                 {
 116                         return NULL;
 117                 }
 118                 return NULL;
 119         }
 120
 121
 122         virtual void OnRawSocketAccept(int fd, const std::string &ip, int localport)
 123         {
 124                 izip_session* session = &sessions[fd];
 125
 126                 /* allocate deflate state */
 127                 session->fd = fd;
 128                 session->status = IZIP_WAITFIRST;
 129
 130                 session->need_bytes = 0;
 131
 132                 session->c_stream.zalloc = (alloc_func)0;
 133                 session->c_stream.zfree = (free_func)0;
 134                 session->c_stream.opaque = (voidpf)0;
 135
 136                 session->d_stream.zalloc = (alloc_func)0;
 137                 session->d_stream.zfree = (free_func)0;
 138                 session->d_stream.opaque = (voidpf)0;
 139
 140         }
 141
 142         virtual void OnRawSocketConnect(int fd)
 143         {
 144                 OnRawSocketAccept(fd, "", 0);
 145         }
 146
 147         virtual void OnRawSocketClose(int fd)
 148         {
 149                 CloseSession(&sessions[fd]);
 150         }
 151
 152         virtual int OnRawSocketRead(int fd, char* buffer, unsigned int count, int &readresult)
 153         {
 154                 izip_session* session = &sessions[fd];
 155
 156                 if (session->status == IZIP_CLOSED)
 157                         return 1;
 158
 159                 int size = 0;
 160
 161                 if (session->need_bytes)
 162                 {
 163                         size = session->need_bytes;
 164                 }
 165                 else
 166                 {
 167                         if (read(fd, &size, sizeof(size)) != sizeof(size))
 168                                 return 0;
 169                         size = ntohl(size);
 170                 }
 171
 172                 ServerInstance->Log(DEBUG,"Size of frame to read: %d%s", size, session->need_bytes ? " (remainder of last frame)" : "");
 173
 174                 unsigned char compr[size+1+session->need_bytes];
 175
 176                 readresult = read(fd, compr + session->need_bytes, size);
 177
 178                 if (readresult == size)
 179                 {
 180                         if(session->status == IZIP_WAITFIRST)
 181                         {
 182                                 session->status = IZIP_OPEN;
 183                         }
 184
 185                         /* Reassemble first part of last frame */
 186                         if (session->need_bytes)
 187                         {
 188                                 for (size_t i = 0; i < session->inbuf.length(); i++)
 189                                         compr[i] = session->inbuf[i];
 190                         }
 191
 192                         session->d_stream.next_in  = (Bytef*)compr;
 193                         session->d_stream.avail_in = 0;
 194                         session->d_stream.next_out = (Bytef*)buffer;
 195                         if (inflateInit(&session->d_stream) != Z_OK)
 196                                 return -EBADF;
 197                         session->status = IZIP_OPEN;
 198
 199                         while ((session->d_stream.total_out < count) && (session->d_stream.total_in < (unsigned int)readresult))
 200                         {
 201                                 session->d_stream.avail_in = session->d_stream.avail_out = 1; /* force small buffers */
 202                                 if (inflate(&session->d_stream, Z_NO_FLUSH) == Z_STREAM_END)
 203                                         break;
 204                         }
 205
 206                         inflateEnd(&session->d_stream);
 207
 208                         total_in_compressed += readresult;
 209                         readresult = session->d_stream.total_out;
 210                         total_in_uncompressed += session->d_stream.total_out;
 211
 212                         buffer[readresult] = 0;
 213                         session->need_bytes = 0;
 214                 }
 215                 else
 216                 {
 217                         /* We need to buffer here */
 218                         ServerInstance->Log(DEBUG,"Didnt read whole frame, got %d bytes of %d!", readresult, size);
 219                         session->need_bytes = ((readresult > -1) ? (size - readresult) : (size));
 220                         if (readresult > 0)
 221                         {
 222                                 /* Do it this way because it needs to be binary safe */
 223                                 for (int i = 0; i < readresult; i++)
 224                                         session->inbuf += compr[i];
 225                         }
 226                 }
 227
 228                 return (readresult > 0);
 229         }
 230
 231         virtual int OnRawSocketWrite(int fd, const char* buffer, int count)
 232         {
 233                 int ocount = count;
 234                 if (!count)
 235                 {
 236                         ServerInstance->Log(DEBUG,"Nothing to do!");
 237                         return 1;
 238                 }
 239
 240                 unsigned char compr[count*2+4];
 241
 242                 izip_session* session = &sessions[fd];
 243
 244                 if(session->status == IZIP_WAITFIRST)
 245                 {
 246                         session->status = IZIP_OPEN;
 247                 }
 248
 249                 if (deflateInit(&session->c_stream, Z_BEST_COMPRESSION) != Z_OK)
 250                 {
 251                         ServerInstance->Log(DEBUG,"Deflate init failed");
 252                 }
 253
 254                 if(session->status != IZIP_OPEN)
 255                 {
 256                         ServerInstance->Log(DEBUG,"State not open!");
 257                         CloseSession(session);
 258                         return 0;
 259                 }
 260
 261                 session->c_stream.next_in  = (Bytef*)buffer;
 262                 session->c_stream.next_out = compr+4;
 263
 264                 while ((session->c_stream.total_in < (unsigned int)count) && (session->c_stream.total_out < (unsigned int)count*2))
 265                 {
 266                         session->c_stream.avail_in = session->c_stream.avail_out = 1; /* force small buffers */
 267                         if (deflate(&session->c_stream, Z_NO_FLUSH) != Z_OK)
 268                         {
 269                                 ServerInstance->Log(DEBUG,"Couldnt deflate!");
 270                                 CloseSession(session);
 271                                 return 0;
 272                         }
 273                 }
 274                 /* Finish the stream, still forcing small buffers: */
 275                 for (;;)
 276                 {
 277                         session->c_stream.avail_out = 1;
 278                         if (deflate(&session->c_stream, Z_FINISH) == Z_STREAM_END)
 279                                 break;
 280                 }
 281
 282                 deflateEnd(&session->c_stream);
 283
 284                 total_out_uncompressed += ocount;
 285                 total_out_compressed += session->c_stream.total_out;
 286
 287                 int x = htonl(session->c_stream.total_out);
 288                 /** XXX: We memcpy it onto the start of the buffer like this to save ourselves a write().
 289                  * A memcpy of 4 or so bytes is less expensive and gives the tcp stack more chance of
 290                  * assembling the frame size into the same packet as the compressed frame.
 291                  */
 292                 memcpy(compr, &x, sizeof(x));
 293                 write(fd, compr, session->c_stream.total_out+4);
 294
 295                 return ocount;
 296         }
 297
 298         void CloseSession(izip_session* session)
 299         {
 300                 session->status = IZIP_CLOSED;
 301         }
 302
 303 };
 304
 305 class ModuleZLibFactory : public ModuleFactory
 306 {
 307  public:
 308         ModuleZLibFactory()
 309         {
 310         }
 311
 312         ~ModuleZLibFactory()
 313         {
 314         }
 315
 316         virtual Module * CreateModule(InspIRCd* Me)
 317         {
 318                 return new ModuleZLib(Me);
 319         }
 320 };
 321
 322
 323 extern "C" void * init_module( void )
 324 {
 325         return new ModuleZLibFactory;
 326 }