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_OPEN, IZIP_CLOSED };
  45
  46 const unsigned int CHUNK = 16384;
  47
  48 class CountedBuffer : public classbase
  49 {
  50         int bufptr;            /* Current tail location */
  51         unsigned char* buffer; /* Current buffer contents */
  52         int bufsz;             /* Current buffer size */
  53         int amount_expected;   /* Amount of data expected */
  54         int amount_read;       /* Amount of data read so far */
  55  public:
  56         CountedBuffer()
  57         {
  58                 bufsz = 1024;
  59                 buffer = new unsigned char[bufsz + 1];
  60                 bufptr = 0;
  61                 amount_read = 0;
  62         }
  63
  64         ~CountedBuffer()
  65         {
  66                 delete[] buffer;
  67         }
  68
  69         void AddData(unsigned char* data, int data_length)
  70         {
  71                 if ((data_length + bufptr) > bufsz)
  72                 {
  73                         /* Buffer is too small, enlarge it and copy contents */
  74                         int old_bufsz = bufsz;
  75                         unsigned char* temp = buffer;
  76
  77                         bufsz += data_length;
  78                         buffer = new unsigned char[bufsz + 1];
  79
  80                         memcpy(buffer, temp, old_bufsz);
  81
  82                         delete[] temp;
  83                 }
  84
  85                 memcpy(buffer + bufptr, data, data_length);
  86                 bufptr += data_length;
  87                 amount_read += data_length;
  88
  89                 if ((!amount_expected) && (amount_read >= 4))
  90                 {
  91                         /* We have enough to read an int */
  92                         int* size = (int*)buffer;
  93                         amount_expected = ntohl(*size);
  94                 }
  95         }
  96
  97         int GetFrame(unsigned char* frame, int maxsize)
  98         {
  99                 if (amount_expected)
 100                 {
 101                         /* We know how much we're expecting...
 102                          * Do we have enough yet?
 103                          */
 104                         if ((amount_read - 4) >= amount_expected)
 105                         {
 106                                 int amt_ex = amount_expected;
 107                                 /* Yes, we have enough now */
 108                                 memcpy(frame, buffer + 4, amount_expected > maxsize ? maxsize : amount_expected);
 109                                 RemoveFirstFrame();
 110                                 return (amt_ex > maxsize) ? maxsize : amt_ex;
 111                         }
 112                 }
 113                 /* Not enough for a frame yet, COME AGAIN! */
 114                 return 0;
 115         }
 116
 117         void RemoveFirstFrame()
 118         {
 119                 unsigned char* temp = buffer;
 120
 121                 bufsz -= (amount_expected + 4);
 122                 buffer = new unsigned char[bufsz + 1];
 123
 124                 memcpy(buffer, temp + amount_expected, bufsz);
 125
 126                 amount_read -= (amount_expected + 4);
 127
 128                 if (amount_read >= 4)
 129                 {
 130                         /* We have enough to read an int */
 131                         int* size = (int*)buffer;
 132                         amount_expected = ntohl(*size);
 133                 }
 134                 else
 135                         amount_expected = 0;
 136
 137                 bufptr = 0;
 138
 139                 delete[] temp;
 140         }
 141 };
 142
 143 /** Represents an ZIP user's extra data
 144  */
 145 class izip_session : public classbase
 146 {
 147  public:
 148         z_stream c_stream; /* compression stream */
 149         z_stream d_stream; /* decompress stream */
 150         izip_status status;
 151         int fd;
 152         CountedBuffer* inbuf;
 153 };
 154
 155 class ModuleZLib : public Module
 156 {
 157         izip_session sessions[MAX_DESCRIPTORS];
 158         float total_out_compressed;
 159         float total_in_compressed;
 160         float total_out_uncompressed;
 161         float total_in_uncompressed;
 162
 163  public:
 164
 165         ModuleZLib(InspIRCd* Me)
 166                 : Module::Module(Me)
 167         {
 168                 ServerInstance->PublishInterface("InspSocketHook", this);
 169
 170                 total_out_compressed = total_in_compressed = 0;
 171                 total_out_uncompressed = total_out_uncompressed = 0;
 172         }
 173
 174         virtual ~ModuleZLib()
 175         {
 176         }
 177
 178         virtual Version GetVersion()
 179         {
 180                 return Version(1, 1, 0, 0, VF_VENDOR, API_VERSION);
 181         }
 182
 183         void Implements(char* List)
 184         {
 185                 List[I_OnRawSocketConnect] = List[I_OnRawSocketAccept] = List[I_OnRawSocketClose] = List[I_OnRawSocketRead] = List[I_OnRawSocketWrite] = 1;
 186                 List[I_OnStats] = List[I_OnRequest] = 1;
 187         }
 188
 189         virtual char* OnRequest(Request* request)
 190         {
 191                 ISHRequest* ISR = (ISHRequest*)request;
 192                 if (strcmp("IS_NAME", request->GetId()) == 0)
 193                 {
 194                         return "zip";
 195                 }
 196                 else if (strcmp("IS_HOOK", request->GetId()) == 0)
 197                 {
 198                         char* ret = "OK";
 199                         try
 200                         {
 201                                 ret = ServerInstance->Config->AddIOHook((Module*)this, (InspSocket*)ISR->Sock) ? (char*)"OK" : NULL;
 202                         }
 203                         catch (ModuleException& e)
 204                         {
 205                                 return NULL;
 206                         }
 207                         return ret;
 208                 }
 209                 else if (strcmp("IS_UNHOOK", request->GetId()) == 0)
 210                 {
 211                         return ServerInstance->Config->DelIOHook((InspSocket*)ISR->Sock) ? (char*)"OK" : NULL;
 212                 }
 213                 else if (strcmp("IS_HSDONE", request->GetId()) == 0)
 214                 {
 215                         return "OK";
 216                 }
 217                 else if (strcmp("IS_ATTACH", request->GetId()) == 0)
 218                 {
 219                         return NULL;
 220                 }
 221                 return NULL;
 222         }
 223
 224         virtual int OnStats(char symbol, userrec* user, string_list &results)
 225         {
 226                 if (symbol == 'z')
 227                 {
 228                         std::string sn = ServerInstance->Config->ServerName;
 229
 230                         float outbound_r = 100 - ((total_out_compressed / (total_out_uncompressed + 0.001)) * 100);
 231                         float inbound_r = 100 - ((total_in_compressed / (total_in_uncompressed + 0.001)) * 100);
 232
 233                         float total_compressed = total_in_compressed + total_out_compressed;
 234                         float total_uncompressed = total_in_uncompressed + total_out_uncompressed;
 235
 236                         float total_r = 100 - ((total_compressed / (total_uncompressed + 0.001)) * 100);
 237
 238                         char outbound_ratio[MAXBUF], inbound_ratio[MAXBUF], combined_ratio[MAXBUF];
 239
 240                         sprintf(outbound_ratio, "%3.2f%%", outbound_r);
 241                         sprintf(inbound_ratio, "%3.2f%%", inbound_r);
 242                         sprintf(combined_ratio, "%3.2f%%", total_r);
 243
 244                         results.push_back(sn+" 304 "+user->nick+" :ZIPSTATS outbound_compressed   = "+ConvToStr(total_out_compressed));
 245                         results.push_back(sn+" 304 "+user->nick+" :ZIPSTATS inbound_compressed    = "+ConvToStr(total_in_compressed));
 246                         results.push_back(sn+" 304 "+user->nick+" :ZIPSTATS outbound_uncompressed = "+ConvToStr(total_out_uncompressed));
 247                         results.push_back(sn+" 304 "+user->nick+" :ZIPSTATS inbound_uncompressed  = "+ConvToStr(total_in_uncompressed));
 248                         results.push_back(sn+" 304 "+user->nick+" :ZIPSTATS outbound_ratio        = "+outbound_ratio);
 249                         results.push_back(sn+" 304 "+user->nick+" :ZIPSTATS inbound_ratio         = "+inbound_ratio);
 250                         results.push_back(sn+" 304 "+user->nick+" :ZIPSTATS combined_ratio        = "+combined_ratio);
 251                         return 0;
 252                 }
 253
 254                 return 0;
 255         }
 256
 257         virtual void OnRawSocketAccept(int fd, const std::string &ip, int localport)
 258         {
 259                 izip_session* session = &sessions[fd];
 260
 261                 /* allocate deflate state */
 262                 session->fd = fd;
 263                 session->status = IZIP_OPEN;
 264
 265                 session->inbuf = new CountedBuffer();
 266                 ServerInstance->Log(DEBUG,"session->inbuf ALLOC = %d, %08x", fd, session->inbuf);
 267
 268                 session->c_stream.zalloc = (alloc_func)0;
 269                 session->c_stream.zfree = (free_func)0;
 270                 session->c_stream.opaque = (voidpf)0;
 271
 272                 session->d_stream.zalloc = (alloc_func)0;
 273                 session->d_stream.zfree = (free_func)0;
 274                 session->d_stream.opaque = (voidpf)0;
 275         }
 276
 277         virtual void OnRawSocketConnect(int fd)
 278         {
 279                 OnRawSocketAccept(fd, "", 0);
 280         }
 281
 282         virtual void OnRawSocketClose(int fd)
 283         {
 284                 CloseSession(&sessions[fd]);
 285         }
 286
 287         virtual int OnRawSocketRead(int fd, char* buffer, unsigned int count, int &readresult)
 288         {
 289                 izip_session* session = &sessions[fd];
 290
 291                 if (session->status == IZIP_CLOSED)
 292                         return 1;
 293
 294                 unsigned char compr[CHUNK + 1];
 295
 296                 readresult = read(fd, compr, CHUNK);
 297
 298                 if (readresult > 0)
 299                 {
 300                         ServerInstance->Log(DEBUG,"session->inbuf PTR = %d, %08x", fd, session->inbuf);
 301
 302                         session->inbuf->AddData(compr, readresult);
 303
 304                         int size = session->inbuf->GetFrame(compr, CHUNK);
 305                         if (size)
 306                         {
 307
 308                                 session->d_stream.next_in  = (Bytef*)compr;
 309                                 session->d_stream.avail_in = 0;
 310                                 session->d_stream.next_out = (Bytef*)buffer;
 311                                 if (inflateInit(&session->d_stream) != Z_OK)
 312                                         return -EBADF;
 313
 314                                 while ((session->d_stream.total_out < count) && (session->d_stream.total_in < (unsigned int)size))
 315                                 {
 316                                         session->d_stream.avail_in = session->d_stream.avail_out = 1;
 317                                         if (inflate(&session->d_stream, Z_NO_FLUSH) == Z_STREAM_END)
 318                                                 break;
 319                                 }
 320
 321                                 inflateEnd(&session->d_stream);
 322
 323                                 total_in_compressed += readresult;
 324                                 readresult = session->d_stream.total_out;
 325                                 total_in_uncompressed += session->d_stream.total_out;
 326
 327                                 buffer[session->d_stream.total_out] = 0;
 328
 329                                 ServerInstance->Log(DEBUG,"Decompressed %d bytes", session->d_stream.total_out);
 330                         }
 331                 }
 332                 return (readresult > 0);
 333         }
 334
 335         virtual int OnRawSocketWrite(int fd, const char* buffer, int count)
 336         {
 337                 ServerInstance->Log(DEBUG,"Compressing %d bytes", count);
 338
 339                 izip_session* session = &sessions[fd];
 340                 int ocount = count;
 341
 342                 if (!count)
 343                 {
 344                         ServerInstance->Log(DEBUG,"Nothing to do!");
 345                         return 1;
 346                 }
 347
 348                 if(session->status != IZIP_OPEN)
 349                 {
 350                         CloseSession(session);
 351                         return 0;
 352                 }
 353
 354                 unsigned char compr[count*2+4];
 355
 356                 if (deflateInit(&session->c_stream, Z_BEST_COMPRESSION) != Z_OK)
 357                 {
 358                         ServerInstance->Log(DEBUG,"Deflate init failed");
 359                 }
 360
 361                 session->c_stream.next_in  = (Bytef*)buffer;
 362                 session->c_stream.next_out = compr+4;
 363
 364                 while ((session->c_stream.total_in < (unsigned int)count) && (session->c_stream.total_out < (unsigned int)count*2))
 365                 {
 366                         session->c_stream.avail_in = session->c_stream.avail_out = 1; /* force small buffers */
 367                         if (deflate(&session->c_stream, Z_NO_FLUSH) != Z_OK)
 368                         {
 369                                 ServerInstance->Log(DEBUG,"Couldnt deflate!");
 370                                 CloseSession(session);
 371                                 return 0;
 372                         }
 373                 }
 374                 /* Finish the stream, still forcing small buffers: */
 375                 for (;;)
 376                 {
 377                         session->c_stream.avail_out = 1;
 378                         if (deflate(&session->c_stream, Z_FINISH) == Z_STREAM_END)
 379                                 break;
 380                 }
 381
 382                 deflateEnd(&session->c_stream);
 383
 384                 total_out_uncompressed += ocount;
 385                 total_out_compressed += session->c_stream.total_out;
 386
 387                 int x = htonl(session->c_stream.total_out);
 388                 /** XXX: We memcpy it onto the start of the buffer like this to save ourselves a write().
 389                  * A memcpy of 4 or so bytes is less expensive and gives the tcp stack more chance of
 390                  * assembling the frame size into the same packet as the compressed frame.
 391                  */
 392                 memcpy(compr, &x, sizeof(x));
 393                 write(fd, compr, session->c_stream.total_out+4);
 394
 395                 return ocount;
 396         }
 397
 398         void CloseSession(izip_session* session)
 399         {
 400                 if (session->status = IZIP_OPEN)
 401                 {
 402                         session->status = IZIP_CLOSED;
 403                         delete session->inbuf;
 404                 }
 405         }
 406
 407 };
 408
 409 class ModuleZLibFactory : public ModuleFactory
 410 {
 411  public:
 412         ModuleZLibFactory()
 413         {
 414         }
 415
 416         ~ModuleZLibFactory()
 417         {
 418         }
 419
 420         virtual Module * CreateModule(InspIRCd* Me)
 421         {
 422                 return new ModuleZLib(Me);
 423         }
 424 };
 425
 426
 427 extern "C" void * init_module( void )
 428 {
 429         return new ModuleZLibFactory;
 430 }