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