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_OnStats] = 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         virtual int OnStats(char symbol, userrec* user, string_list &results)
 122         {
 123                 if (symbol == 'z')
 124                 {
 125                         std::string sn = ServerInstance->Config->ServerName;
 126
 127                         float outbound_r = 100 - ((total_out_compressed / (total_out_uncompressed + 0.001)) * 100);
 128                         float inbound_r = 100 - ((total_in_compressed / (total_in_uncompressed + 0.001)) * 100);
 129
 130                         float total_compressed = total_in_compressed + total_out_compressed;
 131                         float total_uncompressed = total_in_uncompressed + total_out_uncompressed;
 132
 133                         float total_r = 100 - ((total_compressed / (total_uncompressed + 0.001)) * 100);
 134
 135                         char outbound_ratio[MAXBUF], inbound_ratio[MAXBUF], combined_ratio[MAXBUF];
 136
 137                         sprintf(outbound_ratio, "%3.2f%%", outbound_r);
 138                         sprintf(inbound_ratio, "%3.2f%%", inbound_r);
 139                         sprintf(combined_ratio, "%3.2f%%", total_r);
 140
 141                         results.push_back(sn+" 304 "+user->nick+" :ZIPSTATS outbound_compressed   = "+ConvToStr(total_out_compressed));
 142                         results.push_back(sn+" 304 "+user->nick+" :ZIPSTATS inbound_compressed    = "+ConvToStr(total_in_compressed));
 143                         results.push_back(sn+" 304 "+user->nick+" :ZIPSTATS outbound_uncompressed = "+ConvToStr(total_out_uncompressed));
 144                         results.push_back(sn+" 304 "+user->nick+" :ZIPSTATS inbound_uncompressed  = "+ConvToStr(total_in_uncompressed));
 145                         results.push_back(sn+" 304 "+user->nick+" :ZIPSTATS outbound_ratio        = "+outbound_ratio);
 146                         results.push_back(sn+" 304 "+user->nick+" :ZIPSTATS inbound_ratio         = "+inbound_ratio);
 147                         results.push_back(sn+" 304 "+user->nick+" :ZIPSTATS combined_ratio        = "+combined_ratio);
 148                         return 0;
 149                 }
 150
 151                 return 0;
 152         }
 153
 154         virtual void OnRawSocketAccept(int fd, const std::string &ip, int localport)
 155         {
 156                 izip_session* session = &sessions[fd];
 157
 158                 /* allocate deflate state */
 159                 session->fd = fd;
 160                 session->status = IZIP_WAITFIRST;
 161
 162                 session->need_bytes = 0;
 163
 164                 session->c_stream.zalloc = (alloc_func)0;
 165                 session->c_stream.zfree = (free_func)0;
 166                 session->c_stream.opaque = (voidpf)0;
 167
 168                 session->d_stream.zalloc = (alloc_func)0;
 169                 session->d_stream.zfree = (free_func)0;
 170                 session->d_stream.opaque = (voidpf)0;
 171
 172         }
 173
 174         virtual void OnRawSocketConnect(int fd)
 175         {
 176                 OnRawSocketAccept(fd, "", 0);
 177         }
 178
 179         virtual void OnRawSocketClose(int fd)
 180         {
 181                 CloseSession(&sessions[fd]);
 182         }
 183
 184         virtual int OnRawSocketRead(int fd, char* buffer, unsigned int count, int &readresult)
 185         {
 186                 izip_session* session = &sessions[fd];
 187
 188                 if (session->status == IZIP_CLOSED)
 189                         return 1;
 190
 191                 int size = 0;
 192
 193                 if (session->need_bytes)
 194                 {
 195                         size = session->need_bytes;
 196                 }
 197                 else
 198                 {
 199                         if (read(fd, &size, sizeof(size)) != sizeof(size))
 200                                 return 0;
 201                         size = ntohl(size);
 202                 }
 203
 204                 ServerInstance->Log(DEBUG,"Size of frame to read: %d%s", size, session->need_bytes ? " (remainder of last frame)" : "");
 205
 206                 unsigned char compr[size+1+session->need_bytes];
 207
 208                 readresult = read(fd, compr + session->need_bytes, size);
 209
 210                 if (readresult == size)
 211                 {
 212                         if(session->status == IZIP_WAITFIRST)
 213                         {
 214                                 session->status = IZIP_OPEN;
 215                         }
 216
 217                         /* Reassemble first part of last frame */
 218                         if (session->need_bytes)
 219                         {
 220                                 for (size_t i = 0; i < session->inbuf.length(); i++)
 221                                         compr[i] = session->inbuf[i];
 222                         }
 223
 224                         session->d_stream.next_in  = (Bytef*)compr;
 225                         session->d_stream.avail_in = 0;
 226                         session->d_stream.next_out = (Bytef*)buffer;
 227                         if (inflateInit(&session->d_stream) != Z_OK)
 228                                 return -EBADF;
 229                         session->status = IZIP_OPEN;
 230
 231                         while ((session->d_stream.total_out < count) && (session->d_stream.total_in < (unsigned int)readresult))
 232                         {
 233                                 session->d_stream.avail_in = session->d_stream.avail_out = 1; /* force small buffers */
 234                                 if (inflate(&session->d_stream, Z_NO_FLUSH) == Z_STREAM_END)
 235                                         break;
 236                         }
 237
 238                         inflateEnd(&session->d_stream);
 239
 240                         total_in_compressed += readresult;
 241                         readresult = session->d_stream.total_out;
 242                         total_in_uncompressed += session->d_stream.total_out;
 243
 244                         buffer[readresult] = 0;
 245                         session->need_bytes = 0;
 246                 }
 247                 else
 248                 {
 249                         /* We need to buffer here */
 250                         ServerInstance->Log(DEBUG,"Didnt read whole frame, got %d bytes of %d!", readresult, size);
 251                         session->need_bytes = ((readresult > -1) ? (size - readresult) : (size));
 252                         if (readresult > 0)
 253                         {
 254                                 /* Do it this way because it needs to be binary safe */
 255                                 for (int i = 0; i < readresult; i++)
 256                                         session->inbuf += compr[i];
 257                         }
 258                 }
 259
 260                 return (readresult > 0);
 261         }
 262
 263         virtual int OnRawSocketWrite(int fd, const char* buffer, int count)
 264         {
 265                 int ocount = count;
 266                 if (!count)
 267                 {
 268                         ServerInstance->Log(DEBUG,"Nothing to do!");
 269                         return 1;
 270                 }
 271
 272                 unsigned char compr[count*2+4];
 273
 274                 izip_session* session = &sessions[fd];
 275
 276                 if(session->status == IZIP_WAITFIRST)
 277                 {
 278                         session->status = IZIP_OPEN;
 279                 }
 280
 281                 if (deflateInit(&session->c_stream, Z_BEST_COMPRESSION) != Z_OK)
 282                 {
 283                         ServerInstance->Log(DEBUG,"Deflate init failed");
 284                 }
 285
 286                 if(session->status != IZIP_OPEN)
 287                 {
 288                         ServerInstance->Log(DEBUG,"State not open!");
 289                         CloseSession(session);
 290                         return 0;
 291                 }
 292
 293                 session->c_stream.next_in  = (Bytef*)buffer;
 294                 session->c_stream.next_out = compr+4;
 295
 296                 while ((session->c_stream.total_in < (unsigned int)count) && (session->c_stream.total_out < (unsigned int)count*2))
 297                 {
 298                         session->c_stream.avail_in = session->c_stream.avail_out = 1; /* force small buffers */
 299                         if (deflate(&session->c_stream, Z_NO_FLUSH) != Z_OK)
 300                         {
 301                                 ServerInstance->Log(DEBUG,"Couldnt deflate!");
 302                                 CloseSession(session);
 303                                 return 0;
 304                         }
 305                 }
 306                 /* Finish the stream, still forcing small buffers: */
 307                 for (;;)
 308                 {
 309                         session->c_stream.avail_out = 1;
 310                         if (deflate(&session->c_stream, Z_FINISH) == Z_STREAM_END)
 311                                 break;
 312                 }
 313
 314                 deflateEnd(&session->c_stream);
 315
 316                 total_out_uncompressed += ocount;
 317                 total_out_compressed += session->c_stream.total_out;
 318
 319                 int x = htonl(session->c_stream.total_out);
 320                 /** XXX: We memcpy it onto the start of the buffer like this to save ourselves a write().
 321                  * A memcpy of 4 or so bytes is less expensive and gives the tcp stack more chance of
 322                  * assembling the frame size into the same packet as the compressed frame.
 323                  */
 324                 memcpy(compr, &x, sizeof(x));
 325                 write(fd, compr, session->c_stream.total_out+4);
 326
 327                 return ocount;
 328         }
 329
 330         void CloseSession(izip_session* session)
 331         {
 332                 session->status = IZIP_CLOSED;
 333         }
 334
 335 };
 336
 337 class ModuleZLibFactory : public ModuleFactory
 338 {
 339  public:
 340         ModuleZLibFactory()
 341         {
 342         }
 343
 344         ~ModuleZLibFactory()
 345         {
 346         }
 347
 348         virtual Module * CreateModule(InspIRCd* Me)
 349         {
 350                 return new ModuleZLib(Me);
 351         }
 352 };
 353
 354
 355 extern "C" void * init_module( void )
 356 {
 357         return new ModuleZLibFactory;
 358 }