123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156 |
- /*
- * Copyright 2015-present Facebook, Inc.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
- #pragma once
- #include <folly/Optional.h>
- #include <folly/io/async/AsyncSSLSocket.h>
- #include <folly/io/async/AsyncSocket.h>
- #include <folly/io/async/SSLContext.h>
- class BlockingSocket : public folly::AsyncSocket::ConnectCallback,
- public folly::AsyncTransportWrapper::ReadCallback,
- public folly::AsyncTransportWrapper::WriteCallback {
- public:
- explicit BlockingSocket(int fd)
- : sock_(new folly::AsyncSocket(&eventBase_, fd)) {}
- BlockingSocket(
- folly::SocketAddress address,
- std::shared_ptr<folly::SSLContext> sslContext)
- : sock_(
- sslContext ? new folly::AsyncSSLSocket(sslContext, &eventBase_)
- : new folly::AsyncSocket(&eventBase_)),
- address_(address) {}
- explicit BlockingSocket(folly::AsyncSocket::UniquePtr socket)
- : sock_(std::move(socket)) {
- sock_->attachEventBase(&eventBase_);
- }
- void enableTFO() {
- sock_->enableTFO();
- }
- void setAddress(folly::SocketAddress address) {
- address_ = address;
- }
- void open(
- std::chrono::milliseconds timeout = std::chrono::milliseconds::zero()) {
- sock_->connect(this, address_, timeout.count());
- eventBase_.loop();
- if (err_.hasValue()) {
- throw err_.value();
- }
- }
- void close() {
- sock_->close();
- }
- void closeWithReset() {
- sock_->closeWithReset();
- }
- int32_t write(uint8_t const* buf, size_t len) {
- sock_->write(this, buf, len);
- eventBase_.loop();
- if (err_.hasValue()) {
- throw err_.value();
- }
- return len;
- }
- void flush() {}
- int32_t readAll(uint8_t* buf, size_t len) {
- return readHelper(buf, len, true);
- }
- int32_t read(uint8_t* buf, size_t len) {
- return readHelper(buf, len, false);
- }
- int getSocketFD() const {
- return sock_->getFd();
- }
- folly::AsyncSocket* getSocket() {
- return sock_.get();
- }
- folly::AsyncSSLSocket* getSSLSocket() {
- return dynamic_cast<folly::AsyncSSLSocket*>(sock_.get());
- }
- private:
- folly::EventBase eventBase_;
- folly::AsyncSocket::UniquePtr sock_;
- folly::Optional<folly::AsyncSocketException> err_;
- uint8_t* readBuf_{nullptr};
- size_t readLen_{0};
- folly::SocketAddress address_;
- void connectSuccess() noexcept override {}
- void connectErr(const folly::AsyncSocketException& ex) noexcept override {
- err_ = ex;
- }
- void getReadBuffer(void** bufReturn, size_t* lenReturn) override {
- *bufReturn = readBuf_;
- *lenReturn = readLen_;
- }
- void readDataAvailable(size_t len) noexcept override {
- readBuf_ += len;
- readLen_ -= len;
- if (readLen_ == 0) {
- sock_->setReadCB(nullptr);
- }
- }
- void readEOF() noexcept override {}
- void readErr(const folly::AsyncSocketException& ex) noexcept override {
- err_ = ex;
- }
- void writeSuccess() noexcept override {}
- void writeErr(
- size_t /* bytesWritten */,
- const folly::AsyncSocketException& ex) noexcept override {
- err_ = ex;
- }
- int32_t readHelper(uint8_t* buf, size_t len, bool all) {
- if (!sock_->good()) {
- return 0;
- }
- readBuf_ = buf;
- readLen_ = len;
- sock_->setReadCB(this);
- while (!err_ && sock_->good() && readLen_ > 0) {
- eventBase_.loopOnce();
- if (!all) {
- break;
- }
- }
- sock_->setReadCB(nullptr);
- if (err_.hasValue()) {
- throw err_.value();
- }
- if (all && readLen_ > 0) {
- throw folly::AsyncSocketException(
- folly::AsyncSocketException::UNKNOWN, "eof");
- }
- return len - readLen_;
- }
- };
|