folly/io/async/test/BlockingSocket.h

   1 /*
   2  * Copyright 2016 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/AsyncSSLSocket.h>
  20 #include <folly/io/async/AsyncSocket.h>
  21 #include <folly/io/async/SSLContext.h>
  22
  23 class BlockingSocket : public folly::AsyncSocket::ConnectCallback,
  24                        public folly::AsyncTransportWrapper::ReadCallback,
  25                        public folly::AsyncTransportWrapper::WriteCallback {
  26  public:
  27   explicit BlockingSocket(int fd)
  28       : sock_(new folly::AsyncSocket(&eventBase_, fd)) {}
  29
  30   BlockingSocket(
  31       folly::SocketAddress address,
  32       std::shared_ptr<folly::SSLContext> sslContext)
  33       : sock_(
  34             sslContext ? new folly::AsyncSSLSocket(sslContext, &eventBase_)
  35                        : new folly::AsyncSocket(&eventBase_)),
  36         address_(address) {}
  37
  38   explicit BlockingSocket(folly::AsyncSocket::UniquePtr socket)
  39       : sock_(std::move(socket)) {
  40     sock_->attachEventBase(&eventBase_);
  41   }
  42
  43   void enableTFO() {
  44     sock_->enableTFO();
  45   }
  46
  47   void setAddress(folly::SocketAddress address) {
  48     address_ = address;
  49   }
  50
  51   void open(
  52       std::chrono::milliseconds timeout = std::chrono::milliseconds::zero()) {
  53     sock_->connect(this, address_, timeout.count());
  54     eventBase_.loop();
  55     if (err_.hasValue()) {
  56       throw err_.value();
  57     }
  58   }
  59   void close() {
  60     sock_->close();
  61   }
  62   void closeWithReset() {
  63     sock_->closeWithReset();
  64   }
  65
  66   int32_t write(uint8_t const* buf, size_t len) {
  67     sock_->write(this, buf, len);
  68     eventBase_.loop();
  69     if (err_.hasValue()) {
  70       throw err_.value();
  71     }
  72     return len;
  73   }
  74
  75   void flush() {}
  76
  77   int32_t readAll(uint8_t* buf, size_t len) {
  78     return readHelper(buf, len, true);
  79   }
  80
  81   int32_t read(uint8_t* buf, size_t len) {
  82     return readHelper(buf, len, false);
  83   }
  84
  85   int getSocketFD() const {
  86     return sock_->getFd();
  87   }
  88
  89  private:
  90   folly::EventBase eventBase_;
  91   folly::AsyncSocket::UniquePtr sock_;
  92   folly::Optional<folly::AsyncSocketException> err_;
  93   uint8_t* readBuf_{nullptr};
  94   size_t readLen_{0};
  95   folly::SocketAddress address_;
  96
  97   void connectSuccess() noexcept override {}
  98   void connectErr(const folly::AsyncSocketException& ex) noexcept override {
  99     err_ = ex;
 100   }
 101   void getReadBuffer(void** bufReturn, size_t* lenReturn) override {
 102     *bufReturn = readBuf_;
 103     *lenReturn = readLen_;
 104   }
 105   void readDataAvailable(size_t len) noexcept override {
 106     readBuf_ += len;
 107     readLen_ -= len;
 108     if (readLen_ == 0) {
 109       sock_->setReadCB(nullptr);
 110     }
 111   }
 112   void readEOF() noexcept override {}
 113   void readErr(const folly::AsyncSocketException& ex) noexcept override {
 114     err_ = ex;
 115   }
 116   void writeSuccess() noexcept override {}
 117   void writeErr(
 118       size_t /* bytesWritten */,
 119       const folly::AsyncSocketException& ex) noexcept override {
 120     err_ = ex;
 121   }
 122
 123   int32_t readHelper(uint8_t* buf, size_t len, bool all) {
 124     if (!sock_->good()) {
 125       return 0;
 126     }
 127
 128     readBuf_ = buf;
 129     readLen_ = len;
 130     sock_->setReadCB(this);
 131     while (!err_ && sock_->good() && readLen_ > 0) {
 132       eventBase_.loopOnce();
 133       if (!all) {
 134         break;
 135       }
 136     }
 137     sock_->setReadCB(nullptr);
 138     if (err_.hasValue()) {
 139       throw err_.value();
 140     }
 141     if (all && readLen_ > 0) {
 142       throw folly::AsyncSocketException(
 143           folly::AsyncSocketException::UNKNOWN, "eof");
 144     }
 145     return len - readLen_;
 146   }
 147 };