folly/io/async/test/BlockingSocket.h

   1 /*
   2  * Copyright 2015 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16 #pragma once
  17
  18 #include <folly/Optional.h>
  19 #include <folly/io/async/SSLContext.h>
  20 #include <folly/io/async/AsyncSocket.h>
  21 #include <folly/io/async/AsyncSSLSocket.h>
  22
  23 class BlockingSocket : public folly::AsyncSocket::ConnectCallback,
  24                        public folly::AsyncTransportWrapper::ReadCallback,
  25                        public folly::AsyncTransportWrapper::WriteCallback
  26 {
  27  public:
  28   explicit BlockingSocket(int fd)
  29     : sock_(new folly::AsyncSocket(&eventBase_, fd)) {
  30   }
  31
  32   BlockingSocket(folly::SocketAddress address,
  33                  std::shared_ptr<folly::SSLContext> sslContext)
  34     : sock_(sslContext ? new folly::AsyncSSLSocket(sslContext, &eventBase_) :
  35             new folly::AsyncSocket(&eventBase_)),
  36     address_(address) {}
  37
  38   void open() {
  39     sock_->connect(this, address_);
  40     eventBase_.loop();
  41     if (err_.hasValue()) {
  42       throw err_.value();
  43     }
  44   }
  45   void close() {
  46     sock_->close();
  47   }
  48
  49   int32_t write(uint8_t const* buf, size_t len) {
  50     sock_->write(this, buf, len);
  51     eventBase_.loop();
  52     if (err_.hasValue()) {
  53       throw err_.value();
  54     }
  55     return len;
  56   }
  57
  58   void flush() {}
  59
  60   int32_t readAll(uint8_t *buf, size_t len) {
  61     return readHelper(buf, len, true);
  62   }
  63
  64   int32_t read(uint8_t *buf, size_t len) {
  65     return readHelper(buf, len, false);
  66   }
  67
  68   int getSocketFD() const {
  69     return sock_->getFd();
  70   }
  71
  72  private:
  73   folly::EventBase eventBase_;
  74   folly::AsyncSocket::UniquePtr sock_;
  75   folly::Optional<folly::AsyncSocketException> err_;
  76   uint8_t *readBuf_{nullptr};
  77   size_t readLen_{0};
  78   folly::SocketAddress address_;
  79
  80   void connectSuccess() noexcept override {}
  81   void connectErr(const folly::AsyncSocketException& ex) noexcept override {
  82     err_ = ex;
  83   }
  84   void getReadBuffer(void** bufReturn, size_t* lenReturn) override {
  85     *bufReturn = readBuf_;
  86     *lenReturn = readLen_;
  87   }
  88   void readDataAvailable(size_t len) noexcept override {
  89     readBuf_ += len;
  90     readLen_ -= len;
  91     if (readLen_ == 0) {
  92       sock_->setReadCB(nullptr);
  93     }
  94   }
  95   void readEOF() noexcept override {
  96   }
  97   void readErr(const folly::AsyncSocketException& ex) noexcept override {
  98     err_ = ex;
  99   }
 100   void writeSuccess() noexcept override {}
 101   void writeErr(size_t /* bytesWritten */,
 102                 const folly::AsyncSocketException& ex) noexcept override {
 103     err_ = ex;
 104   }
 105
 106   int32_t readHelper(uint8_t *buf, size_t len, bool all) {
 107     readBuf_ = buf;
 108     readLen_ = len;
 109     sock_->setReadCB(this);
 110     while (!err_ && sock_->good() && readLen_ > 0) {
 111       eventBase_.loop();
 112       if (!all) {
 113         break;
 114       }
 115     }
 116     sock_->setReadCB(nullptr);
 117     if (err_.hasValue()) {
 118       throw err_.value();
 119     }
 120     if (all && readLen_ > 0) {
 121       throw folly::AsyncSocketException(folly::AsyncSocketException::UNKNOWN,
 122                                         "eof");
 123     }
 124     return len - readLen_;
 125   }
 126 };