lib/cpp/src/transport/TTransportUtils.h

   1 // Copyright (c) 2006- Facebook
   2 // Distributed under the Thrift Software License
   3 //
   4 // See accompanying file LICENSE or visit the Thrift site at:
   5 // http://developers.facebook.com/thrift/
   6
   7 #ifndef _THRIFT_TRANSPORT_TTRANSPORTUTILS_H_
   8 #define _THRIFT_TRANSPORT_TTRANSPORTUTILS_H_ 1
   9
  10 #include <string>
  11 #include <algorithm>
  12 #include <transport/TTransport.h>
  13 #include <transport/TFileTransport.h>
  14
  15 namespace facebook { namespace thrift { namespace transport {
  16
  17 /**
  18  * The null transport is a dummy transport that doesn't actually do anything.
  19  * It's sort of an analogy to /dev/null, you can never read anything from it
  20  * and it will let you write anything you want to it, though it won't actually
  21  * go anywhere.
  22  *
  23  * @author Mark Slee <mcslee@facebook.com>
  24  */
  25 class TNullTransport : public TTransport {
  26  public:
  27   TNullTransport() {}
  28
  29   ~TNullTransport() {}
  30
  31   bool isOpen() {
  32     return true;
  33   }
  34
  35   void open() {}
  36
  37   void write(const uint8_t* buf, uint32_t len) {
  38     return;
  39   }
  40
  41 };
  42
  43
  44 /**
  45  * Buffered transport. For reads it will read more data than is requested
  46  * and will serve future data out of a local buffer. For writes, data is
  47  * stored to an in memory buffer before being written out.
  48  *
  49  * @author Mark Slee <mcslee@facebook.com>
  50  */
  51 class TBufferedTransport : public TTransport {
  52  public:
  53   TBufferedTransport(boost::shared_ptr<TTransport> transport) :
  54     transport_(transport),
  55     rBufSize_(512), rPos_(0), rLen_(0),
  56     wBufSize_(512), wLen_(0) {
  57     rBuf_ = new uint8_t[rBufSize_];
  58     wBuf_ = new uint8_t[wBufSize_];
  59   }
  60
  61   TBufferedTransport(boost::shared_ptr<TTransport> transport, uint32_t sz) :
  62     transport_(transport),
  63     rBufSize_(sz), rPos_(0), rLen_(0),
  64     wBufSize_(sz), wLen_(0) {
  65     rBuf_ = new uint8_t[rBufSize_];
  66     wBuf_ = new uint8_t[wBufSize_];
  67   }
  68
  69   TBufferedTransport(boost::shared_ptr<TTransport> transport, uint32_t rsz, uint32_t wsz) :
  70     transport_(transport),
  71     rBufSize_(rsz), rPos_(0), rLen_(0),
  72     wBufSize_(wsz), wLen_(0) {
  73     rBuf_ = new uint8_t[rBufSize_];
  74     wBuf_ = new uint8_t[wBufSize_];
  75   }
  76
  77   ~TBufferedTransport() {
  78     delete [] rBuf_;
  79     delete [] wBuf_;
  80   }
  81
  82   bool isOpen() {
  83     return transport_->isOpen();
  84   }
  85
  86   bool peek() {
  87     if (rPos_ >= rLen_) {
  88       rLen_ = transport_->read(rBuf_, rBufSize_);
  89       rPos_ = 0;
  90     }
  91     return (rLen_ > rPos_);
  92   }
  93
  94   void open() {
  95     transport_->open();
  96   }
  97
  98   void close() {
  99     flush();
 100     transport_->close();
 101   }
 102
 103   uint32_t read(uint8_t* buf, uint32_t len);
 104
 105   void write(const uint8_t* buf, uint32_t len);
 106
 107   void flush();
 108
 109   bool borrow(uint8_t* buf, uint32_t len);
 110
 111   void consume(uint32_t len);
 112
 113  protected:
 114   boost::shared_ptr<TTransport> transport_;
 115   uint8_t* rBuf_;
 116   uint32_t rBufSize_;
 117   uint32_t rPos_;
 118   uint32_t rLen_;
 119
 120   uint8_t* wBuf_;
 121   uint32_t wBufSize_;
 122   uint32_t wLen_;
 123 };
 124
 125 /**
 126  * Wraps a transport into a buffered one.
 127  *
 128  * @author Mark Slee <mcslee@facebook.com>
 129  */
 130 class TBufferedTransportFactory : public TTransportFactory {
 131  public:
 132   TBufferedTransportFactory() {}
 133
 134   virtual ~TBufferedTransportFactory() {}
 135
 136   /**
 137    * Wraps the transport into a buffered one.
 138    */
 139   virtual boost::shared_ptr<TTransport> getTransport(boost::shared_ptr<TTransport> trans) {
 140     return boost::shared_ptr<TTransport>(new TBufferedTransport(trans));
 141   }
 142
 143 };
 144
 145 /**
 146  * Framed transport. All writes go into an in-memory buffer until flush is
 147  * called, at which point the transport writes the length of the entire
 148  * binary chunk followed by the data payload. This allows the receiver on the
 149  * other end to always do fixed-length reads.
 150  *
 151  * @author Mark Slee <mcslee@facebook.com>
 152  */
 153 class TFramedTransport : public TTransport {
 154  public:
 155   TFramedTransport(boost::shared_ptr<TTransport> transport) :
 156     transport_(transport),
 157     rPos_(0),
 158     rLen_(0),
 159     read_(true),
 160     wBufSize_(512),
 161     wLen_(0),
 162     write_(true) {
 163     rBuf_ = NULL;
 164     wBuf_ = new uint8_t[wBufSize_];
 165   }
 166
 167   TFramedTransport(boost::shared_ptr<TTransport> transport, uint32_t sz) :
 168     transport_(transport),
 169     rPos_(0),
 170     rLen_(0),
 171     read_(true),
 172     wBufSize_(sz),
 173     wLen_(0),
 174     write_(true) {
 175     rBuf_ = NULL;
 176     wBuf_ = new uint8_t[wBufSize_];
 177   }
 178
 179   ~TFramedTransport() {
 180     if (rBuf_ != NULL) {
 181       delete [] rBuf_;
 182     }
 183     if (wBuf_ != NULL) {
 184       delete [] wBuf_;
 185     }
 186   }
 187
 188   void setRead(bool read) {
 189     read_ = read;
 190   }
 191
 192   void setWrite(bool write) {
 193     write_ = write;
 194   }
 195
 196   void open() {
 197     transport_->open();
 198   }
 199
 200   bool isOpen() {
 201     return transport_->isOpen();
 202   }
 203
 204   bool peek() {
 205     if (rPos_ < rLen_) {
 206       return true;
 207     }
 208     return transport_->peek();
 209   }
 210
 211   void close() {
 212     if (wLen_ > 0) {
 213       flush();
 214     }
 215     transport_->close();
 216   }
 217
 218   uint32_t read(uint8_t* buf, uint32_t len);
 219
 220   void write(const uint8_t* buf, uint32_t len);
 221
 222   void flush();
 223
 224   bool borrow(uint8_t* buf, uint32_t len);
 225
 226   void consume(uint32_t len);
 227
 228  protected:
 229   boost::shared_ptr<TTransport> transport_;
 230   uint8_t* rBuf_;
 231   uint32_t rPos_;
 232   uint32_t rLen_;
 233   bool read_;
 234
 235   uint8_t* wBuf_;
 236   uint32_t wBufSize_;
 237   uint32_t wLen_;
 238   bool write_;
 239
 240   /**
 241    * Reads a frame of input from the underlying stream.
 242    */
 243   void readFrame();
 244 };
 245
 246 /**
 247  * A memory buffer is a tranpsort that simply reads from and writes to an
 248  * in memory buffer. Anytime you call write on it, the data is simply placed
 249  * into a buffer, and anytime you call read, data is read from that buffer.
 250  *
 251  * The buffers are allocated using C constructs malloc,realloc, and the size
 252  * doubles as necessary.
 253  *
 254  * @author Mark Slee <mcslee@facebook.com>
 255  */
 256 class TMemoryBuffer : public TTransport {
 257  private:
 258
 259   // Common initialization done by all constructors.
 260   void initCommon(uint8_t* buf, uint32_t size, bool owner, uint32_t wPos) {
 261     if (buf == NULL && size != 0) {
 262       assert(owner);
 263       buf = (uint8_t*)malloc(size);
 264       if (buf == NULL) {
 265         throw TTransportException("Out of memory");
 266       }
 267     }
 268
 269     buffer_ = buf;
 270     bufferSize_ = size;
 271     owner_ = owner;
 272     wPos_ = wPos;
 273     rPos_ = 0;
 274   }
 275
 276  public:
 277   static const uint32_t defaultSize = 1024;
 278
 279   TMemoryBuffer() {
 280     initCommon(NULL, defaultSize, true, 0);
 281   }
 282
 283   TMemoryBuffer(uint32_t sz) {
 284     initCommon(NULL, sz, true, 0);
 285   }
 286
 287   // transferOwnership should be true if you want TMemoryBuffer to call free(buf).
 288   TMemoryBuffer(uint8_t* buf, int sz, bool transferOwnership = false) {
 289     initCommon(buf, sz, transferOwnership, sz);
 290   }
 291
 292   // copy should be true if you want TMemoryBuffer to make a copy of the string.
 293   // If you set copy to false, the string must not be destroyed before you are
 294   // done with the TMemoryBuffer.
 295   TMemoryBuffer(const std::string& str, bool copy = false) {
 296     if (copy) {
 297       initCommon(NULL, str.length(), true, 0);
 298       this->write(reinterpret_cast<const uint8_t*>(str.data()), str.length());
 299     } else {
 300       // This first argument should be equivalent to the following:
 301       // const_cast<uint8_t*>(reinterpret_cast<const uint8_t*>(str.data()))
 302       initCommon((uint8_t*)str.data(), str.length(), false, str.length());
 303     }
 304   }
 305
 306   ~TMemoryBuffer() {
 307     if (owner_) {
 308       free(buffer_);
 309       buffer_ = NULL;
 310     }
 311   }
 312
 313   bool isOpen() {
 314     return true;
 315   }
 316
 317   bool peek() {
 318     return (rPos_ < wPos_);
 319   }
 320
 321   void open() {}
 322
 323   void close() {}
 324
 325   void getBuffer(uint8_t** bufPtr, uint32_t* sz) {
 326     *bufPtr = buffer_;
 327     *sz = wPos_;
 328   }
 329
 330   std::string getBufferAsString() {
 331     if (buffer_ == NULL) {
 332       return "";
 333     }
 334     return std::string((char*)buffer_, (std::string::size_type)wPos_);
 335   }
 336
 337   void appendBufferToString(std::string& str) {
 338     if (buffer_ == NULL) {
 339       return;
 340     }
 341     str.append((char*)buffer_, wPos_);
 342   }
 343
 344   void resetBuffer() {
 345     wPos_ = 0;
 346     rPos_ = 0;
 347   }
 348
 349   // transferOwnership should be true if you want TMemoryBuffer to call free(buf).
 350   void resetBuffer(uint8_t* buf, uint32_t sz, bool transferOwnership = false) {
 351     if (owner_) {
 352       free(buffer_);
 353     }
 354     owner_ = transferOwnership;
 355     buffer_ = buf;
 356     bufferSize_ = sz;
 357     wPos_ = sz;
 358     rPos_ = 0;
 359   }
 360
 361   // See the constructor that takes a string.
 362   void resetFromString(const std::string& str, bool copy = false) {
 363     if (copy) {
 364       uint8_t* buf = (uint8_t*)malloc(str.length());
 365       if (buf == NULL) {
 366         throw TTransportException("Out of memory");
 367       }
 368       std::copy(str.begin(), str.end(), buf);
 369       resetBuffer(buf, str.length(), true);
 370     } else {
 371       // See the above comment about const_cast.
 372       resetBuffer((uint8_t*)str.data(), str.length(), false);
 373     }
 374   }
 375
 376   uint32_t read(uint8_t* buf, uint32_t len);
 377
 378   std::string readAsString(uint32_t len) {
 379     std::string str;
 380     (void)readAppendToString(str, len);
 381     return str;
 382   }
 383
 384   uint32_t readAppendToString(std::string& str, uint32_t len);
 385
 386   void readEnd() {
 387     if (rPos_ == wPos_) {
 388       resetBuffer();
 389     }
 390   }
 391
 392   void write(const uint8_t* buf, uint32_t len);
 393
 394   uint32_t available() {
 395     return wPos_ - rPos_;
 396   }
 397
 398   bool borrow(uint8_t* buf, uint32_t len);
 399
 400   void consume(uint32_t len);
 401
 402  private:
 403   // Data buffer
 404   uint8_t* buffer_;
 405
 406   // Allocated buffer size
 407   uint32_t bufferSize_;
 408
 409   // Where the write is at
 410   uint32_t wPos_;
 411
 412   // Where the reader is at
 413   uint32_t rPos_;
 414
 415   // Is this object the owner of the buffer?
 416   bool owner_;
 417
 418 };
 419
 420 /**
 421  * TPipedTransport. This transport allows piping of a request from one
 422  * transport to another either when readEnd() or writeEnd(). The typical
 423  * use case for this is to log a request or a reply to disk.
 424  * The underlying buffer expands to a keep a copy of the entire
 425  * request/response.
 426  *
 427  * @author Aditya Agarwal <aditya@facebook.com>
 428  */
 429 class TPipedTransport : virtual public TTransport {
 430  public:
 431   TPipedTransport(boost::shared_ptr<TTransport> srcTrans,
 432                   boost::shared_ptr<TTransport> dstTrans) :
 433     srcTrans_(srcTrans),
 434     dstTrans_(dstTrans),
 435     rBufSize_(512), rPos_(0), rLen_(0),
 436     wBufSize_(512), wLen_(0) {
 437
 438     // default is to to pipe the request when readEnd() is called
 439     pipeOnRead_ = true;
 440     pipeOnWrite_ = false;
 441
 442     rBuf_ = (uint8_t*) malloc(sizeof(uint8_t) * rBufSize_);
 443     wBuf_ = (uint8_t*) malloc(sizeof(uint8_t) * wBufSize_);
 444   }
 445
 446   TPipedTransport(boost::shared_ptr<TTransport> srcTrans,
 447                   boost::shared_ptr<TTransport> dstTrans,
 448                   uint32_t sz) :
 449     srcTrans_(srcTrans),
 450     dstTrans_(dstTrans),
 451     rBufSize_(512), rPos_(0), rLen_(0),
 452     wBufSize_(sz), wLen_(0) {
 453
 454     rBuf_ = (uint8_t*) malloc(sizeof(uint8_t) * rBufSize_);
 455     wBuf_ = (uint8_t*) malloc(sizeof(uint8_t) * wBufSize_);
 456   }
 457
 458   ~TPipedTransport() {
 459     free(rBuf_);
 460     free(wBuf_);
 461   }
 462
 463   bool isOpen() {
 464     return srcTrans_->isOpen();
 465   }
 466
 467   bool peek() {
 468     if (rPos_ >= rLen_) {
 469       // Double the size of the underlying buffer if it is full
 470       if (rLen_ == rBufSize_) {
 471         rBufSize_ *=2;
 472         rBuf_ = (uint8_t *)realloc(rBuf_, sizeof(uint8_t) * rBufSize_);
 473       }
 474
 475       // try to fill up the buffer
 476       rLen_ += srcTrans_->read(rBuf_+rPos_, rBufSize_ - rPos_);
 477     }
 478     return (rLen_ > rPos_);
 479   }
 480
 481
 482   void open() {
 483     srcTrans_->open();
 484   }
 485
 486   void close() {
 487     srcTrans_->close();
 488   }
 489
 490   void setPipeOnRead(bool pipeVal) {
 491     pipeOnRead_ = pipeVal;
 492   }
 493
 494   void setPipeOnWrite(bool pipeVal) {
 495     pipeOnWrite_ = pipeVal;
 496   }
 497
 498   uint32_t read(uint8_t* buf, uint32_t len);
 499
 500   void readEnd() {
 501
 502     if (pipeOnRead_) {
 503       dstTrans_->write(rBuf_, rLen_);
 504       dstTrans_->flush();
 505     }
 506
 507     srcTrans_->readEnd();
 508
 509     // reset state
 510     rLen_ = 0;
 511     rPos_ = 0;
 512   }
 513
 514   void write(const uint8_t* buf, uint32_t len);
 515
 516   void writeEnd() {
 517     if (pipeOnWrite_) {
 518       dstTrans_->write(wBuf_, wLen_);
 519       dstTrans_->flush();
 520     }
 521   }
 522
 523   void flush();
 524
 525   boost::shared_ptr<TTransport> getTargetTransport() {
 526     return dstTrans_;
 527   }
 528
 529  protected:
 530   boost::shared_ptr<TTransport> srcTrans_;
 531   boost::shared_ptr<TTransport> dstTrans_;
 532
 533   uint8_t* rBuf_;
 534   uint32_t rBufSize_;
 535   uint32_t rPos_;
 536   uint32_t rLen_;
 537
 538   uint8_t* wBuf_;
 539   uint32_t wBufSize_;
 540   uint32_t wLen_;
 541
 542   bool pipeOnRead_;
 543   bool pipeOnWrite_;
 544 };
 545
 546
 547 /**
 548  * Wraps a transport into a pipedTransport instance.
 549  *
 550  * @author Aditya Agarwal <aditya@facebook.com>
 551  */
 552 class TPipedTransportFactory : public TTransportFactory {
 553  public:
 554   TPipedTransportFactory() {}
 555   TPipedTransportFactory(boost::shared_ptr<TTransport> dstTrans) {
 556     initializeTargetTransport(dstTrans);
 557   }
 558   virtual ~TPipedTransportFactory() {}
 559
 560   /**
 561    * Wraps the base transport into a piped transport.
 562    */
 563   virtual boost::shared_ptr<TTransport> getTransport(boost::shared_ptr<TTransport> srcTrans) {
 564     return boost::shared_ptr<TTransport>(new TPipedTransport(srcTrans, dstTrans_));
 565   }
 566
 567   virtual void initializeTargetTransport(boost::shared_ptr<TTransport> dstTrans) {
 568     if (dstTrans_.get() == NULL) {
 569       dstTrans_ = dstTrans;
 570     } else {
 571       throw TException("Target transport already initialized");
 572     }
 573   }
 574
 575  protected:
 576   boost::shared_ptr<TTransport> dstTrans_;
 577 };
 578
 579 /**
 580  * TPipedFileTransport. This is just like a TTransport, except that
 581  * it is a templatized class, so that clients who rely on a specific
 582  * TTransport can still access the original transport.
 583  *
 584  * @author James Wang <jwang@facebook.com>
 585  */
 586 class TPipedFileReaderTransport : public TPipedTransport,
 587                                   public TFileReaderTransport {
 588  public:
 589   TPipedFileReaderTransport(boost::shared_ptr<TFileReaderTransport> srcTrans, boost::shared_ptr<TTransport> dstTrans);
 590
 591   ~TPipedFileReaderTransport();
 592
 593   // TTransport functions
 594   bool isOpen();
 595   bool peek();
 596   void open();
 597   void close();
 598   uint32_t read(uint8_t* buf, uint32_t len);
 599   uint32_t readAll(uint8_t* buf, uint32_t len);
 600   void readEnd();
 601   void write(const uint8_t* buf, uint32_t len);
 602   void writeEnd();
 603   void flush();
 604
 605   // TFileReaderTransport functions
 606   int32_t getReadTimeout();
 607   void setReadTimeout(int32_t readTimeout);
 608   uint32_t getNumChunks();
 609   uint32_t getCurChunk();
 610   void seekToChunk(int32_t chunk);
 611   void seekToEnd();
 612
 613  protected:
 614   // shouldn't be used
 615   TPipedFileReaderTransport();
 616   boost::shared_ptr<TFileReaderTransport> srcTrans_;
 617 };
 618
 619 /**
 620  * Creates a TPipedFileReaderTransport from a filepath and a destination transport
 621  *
 622  * @author James Wang <jwang@facebook.com>
 623  */
 624 class TPipedFileReaderTransportFactory : public TPipedTransportFactory {
 625  public:
 626   TPipedFileReaderTransportFactory() {}
 627   TPipedFileReaderTransportFactory(boost::shared_ptr<TTransport> dstTrans)
 628     : TPipedTransportFactory(dstTrans)
 629   {}
 630   virtual ~TPipedFileReaderTransportFactory() {}
 631
 632   boost::shared_ptr<TTransport> getTransport(boost::shared_ptr<TTransport> srcTrans) {
 633     boost::shared_ptr<TFileReaderTransport> pFileReaderTransport = boost::dynamic_pointer_cast<TFileReaderTransport>(srcTrans);
 634     if (pFileReaderTransport.get() != NULL) {
 635       return getFileReaderTransport(pFileReaderTransport);
 636     } else {
 637       return boost::shared_ptr<TTransport>();
 638     }
 639   }
 640
 641   boost::shared_ptr<TFileReaderTransport> getFileReaderTransport(boost::shared_ptr<TFileReaderTransport> srcTrans) {
 642     return boost::shared_ptr<TFileReaderTransport>(new TPipedFileReaderTransport(srcTrans, dstTrans_));
 643   }
 644 };
 645
 646 }}} // facebook::thrift::transport
 647
 648 #endif // #ifndef _THRIFT_TRANSPORT_TTRANSPORTUTILS_H_