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
  65  public:
  66
  67         ModuleZLib(InspIRCd* Me)
  68                 : Module::Module(Me)
  69         {
  70                 ServerInstance->PublishInterface("InspSocketHook", this);
  71         }
  72
  73         virtual ~ModuleZLib()
  74         {
  75         }
  76
  77         virtual Version GetVersion()
  78         {
  79                 return Version(1, 1, 0, 0, VF_VENDOR, API_VERSION);
  80         }
  81
  82         void Implements(char* List)
  83         {
  84                 List[I_OnRawSocketConnect] = List[I_OnRawSocketAccept] = List[I_OnRawSocketClose] = List[I_OnRawSocketRead] = List[I_OnRawSocketWrite] = 1;
  85                 List[I_OnRequest] = 1;
  86         }
  87
  88         virtual char* OnRequest(Request* request)
  89         {
  90                 ISHRequest* ISR = (ISHRequest*)request;
  91                 if (strcmp("IS_NAME", request->GetId()) == 0)
  92                 {
  93                         return "zip";
  94                 }
  95                 else if (strcmp("IS_HOOK", request->GetId()) == 0)
  96                 {
  97                         return ServerInstance->Config->AddIOHook((Module*)this, (InspSocket*)ISR->Sock) ? (char*)"OK" : NULL;
  98                 }
  99                 else if (strcmp("IS_UNHOOK", request->GetId()) == 0)
 100                 {
 101                         return ServerInstance->Config->DelIOHook((InspSocket*)ISR->Sock) ? (char*)"OK" : NULL;
 102                 }
 103                 else if (strcmp("IS_HSDONE", request->GetId()) == 0)
 104                 {
 105                         return "OK";
 106                 }
 107                 else if (strcmp("IS_ATTACH", request->GetId()) == 0)
 108                 {
 109                         return NULL;
 110                 }
 111                 return NULL;
 112         }
 113
 114
 115         virtual void OnRawSocketAccept(int fd, const std::string &ip, int localport)
 116         {
 117                 izip_session* session = &sessions[fd];
 118
 119                 /* allocate deflate state */
 120                 session->fd = fd;
 121                 session->status = IZIP_WAITFIRST;
 122
 123                 session->need_bytes = 0;
 124
 125                 session->c_stream.zalloc = (alloc_func)0;
 126                 session->c_stream.zfree = (free_func)0;
 127                 session->c_stream.opaque = (voidpf)0;
 128
 129                 session->d_stream.zalloc = (alloc_func)0;
 130                 session->d_stream.zfree = (free_func)0;
 131                 session->d_stream.opaque = (voidpf)0;
 132
 133         }
 134
 135         virtual void OnRawSocketConnect(int fd)
 136         {
 137                 OnRawSocketAccept(fd, "", 0);
 138         }
 139
 140         virtual void OnRawSocketClose(int fd)
 141         {
 142                 CloseSession(&sessions[fd]);
 143         }
 144
 145         virtual int OnRawSocketRead(int fd, char* buffer, unsigned int count, int &readresult)
 146         {
 147                 izip_session* session = &sessions[fd];
 148
 149                 if (session->status == IZIP_CLOSED)
 150                         return 1;
 151
 152                 int size = 0;
 153
 154                 if (session->need_bytes)
 155                 {
 156                         size = session->need_bytes;
 157                 }
 158                 else
 159                 {
 160                         if (read(fd, &size, sizeof(size)) != sizeof(size))
 161                                 return 0;
 162                         size = ntohl(size);
 163                 }
 164
 165                 ServerInstance->Log(DEBUG,"Size of frame to read: %d%s", size, session->need_bytes ? " (remainder of last frame)" : "");
 166
 167                 unsigned char compr[size+1+session->need_bytes];
 168
 169                 readresult = read(fd, compr + session->need_bytes, size);
 170
 171                 if (readresult == size)
 172                 {
 173                         if(session->status == IZIP_WAITFIRST)
 174                         {
 175                                 session->status = IZIP_OPEN;
 176                         }
 177
 178                         /* Reassemble first part of last frame */
 179                         if (session->need_bytes)
 180                         {
 181                                 for (size_t i = 0; i < session->inbuf.length(); i++)
 182                                         compr[i] = session->inbuf[i];
 183                         }
 184
 185                         session->d_stream.next_in  = (Bytef*)compr;
 186                         session->d_stream.avail_in = 0;
 187                         session->d_stream.next_out = (Bytef*)buffer;
 188                         if (inflateInit(&session->d_stream) != Z_OK)
 189                                 return -EBADF;
 190                         session->status = IZIP_OPEN;
 191
 192                         while ((session->d_stream.total_out < count) && (session->d_stream.total_in < (unsigned int)readresult))
 193                         {
 194                                 session->d_stream.avail_in = session->d_stream.avail_out = 1; /* force small buffers */
 195                                 if (inflate(&session->d_stream, Z_NO_FLUSH) == Z_STREAM_END)
 196                                         break;
 197                         }
 198
 199                         inflateEnd(&session->d_stream);
 200
 201                         readresult = session->d_stream.total_out;
 202
 203                         buffer[readresult] = 0;
 204                         session->need_bytes = 0;
 205                 }
 206                 else
 207                 {
 208                         /* We need to buffer here */
 209                         ServerInstance->Log(DEBUG,"Didnt read whole frame, got %d bytes of %d!", readresult, size);
 210                         session->need_bytes = ((readresult > -1) ? (size - readresult) : (size));
 211                         if (readresult > 0)
 212                         {
 213                                 /* Do it this way because it needs to be binary safe */
 214                                 for (int i = 0; i < readresult; i++)
 215                                         session->inbuf += compr[i];
 216                         }
 217                 }
 218
 219                 return (readresult > 0);
 220         }
 221
 222         virtual int OnRawSocketWrite(int fd, const char* buffer, int count)
 223         {
 224                 int ocount = count;
 225                 if (!count)
 226                 {
 227                         ServerInstance->Log(DEBUG,"Nothing to do!");
 228                         return 1;
 229                 }
 230
 231                 unsigned char compr[count*2+4];
 232
 233                 izip_session* session = &sessions[fd];
 234
 235                 if(session->status == IZIP_WAITFIRST)
 236                 {
 237                         session->status = IZIP_OPEN;
 238                 }
 239
 240                 if (deflateInit(&session->c_stream, Z_BEST_COMPRESSION) != Z_OK)
 241                 {
 242                         ServerInstance->Log(DEBUG,"Deflate init failed");
 243                 }
 244
 245                 if(session->status != IZIP_OPEN)
 246                 {
 247                         ServerInstance->Log(DEBUG,"State not open!");
 248                         CloseSession(session);
 249                         return 0;
 250                 }
 251
 252                 session->c_stream.next_in  = (Bytef*)buffer;
 253                 session->c_stream.next_out = compr+4;
 254
 255                 while ((session->c_stream.total_in < (unsigned int)count) && (session->c_stream.total_out < (unsigned int)count*2))
 256                 {
 257                         session->c_stream.avail_in = session->c_stream.avail_out = 1; /* force small buffers */
 258                         if (deflate(&session->c_stream, Z_NO_FLUSH) != Z_OK)
 259                         {
 260                                 ServerInstance->Log(DEBUG,"Couldnt deflate!");
 261                                 CloseSession(session);
 262                                 return 0;
 263                         }
 264                 }
 265                 /* Finish the stream, still forcing small buffers: */
 266                 for (;;)
 267                 {
 268                         session->c_stream.avail_out = 1;
 269                         if (deflate(&session->c_stream, Z_FINISH) == Z_STREAM_END)
 270                                 break;
 271                 }
 272
 273                 int x = htonl(session->c_stream.total_out);
 274                 /** XXX: We memcpy it onto the start of the buffer like this to save ourselves a write().
 275                  * A memcpy of 4 or so bytes is less expensive and gives the tcp stack more chance of
 276                  * assembling the frame size into the same packet as the compressed frame.
 277                  */
 278                 memcpy(compr, &x, sizeof(x));
 279                 write(fd, compr, session->c_stream.total_out+4);
 280
 281                 deflateEnd(&session->c_stream);
 282
 283                 return ocount;
 284         }
 285
 286         void CloseSession(izip_session* session)
 287         {
 288                 session->status = IZIP_CLOSED;
 289         }
 290
 291 };
 292
 293 class ModuleZLibFactory : public ModuleFactory
 294 {
 295  public:
 296         ModuleZLibFactory()
 297         {
 298         }
 299
 300         ~ModuleZLibFactory()
 301         {
 302         }
 303
 304         virtual Module * CreateModule(InspIRCd* Me)
 305         {
 306                 return new ModuleZLib(Me);
 307         }
 308 };
 309
 310
 311 extern "C" void * init_module( void )
 312 {
 313         return new ModuleZLibFactory;
 314 }