2 * Copyright 2017 Facebook, Inc.
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
8 * http://www.apache.org/licenses/LICENSE-2.0
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.
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>
23 class BlockingSocket : public folly::AsyncSocket::ConnectCallback,
24 public folly::AsyncTransportWrapper::ReadCallback,
25 public folly::AsyncTransportWrapper::WriteCallback {
27 explicit BlockingSocket(int fd)
28 : sock_(new folly::AsyncSocket(&eventBase_, fd)) {}
31 folly::SocketAddress address,
32 std::shared_ptr<folly::SSLContext> sslContext)
34 sslContext ? new folly::AsyncSSLSocket(sslContext, &eventBase_)
35 : new folly::AsyncSocket(&eventBase_)),
38 explicit BlockingSocket(folly::AsyncSocket::UniquePtr socket)
39 : sock_(std::move(socket)) {
40 sock_->attachEventBase(&eventBase_);
47 void setAddress(folly::SocketAddress address) {
52 std::chrono::milliseconds timeout = std::chrono::milliseconds::zero()) {
53 sock_->connect(this, address_, timeout.count());
55 if (err_.hasValue()) {
62 void closeWithReset() {
63 sock_->closeWithReset();
66 int32_t write(uint8_t const* buf, size_t len) {
67 sock_->write(this, buf, len);
69 if (err_.hasValue()) {
77 int32_t readAll(uint8_t* buf, size_t len) {
78 return readHelper(buf, len, true);
81 int32_t read(uint8_t* buf, size_t len) {
82 return readHelper(buf, len, false);
85 int getSocketFD() const {
86 return sock_->getFd();
90 folly::EventBase eventBase_;
91 folly::AsyncSocket::UniquePtr sock_;
92 folly::Optional<folly::AsyncSocketException> err_;
93 uint8_t* readBuf_{nullptr};
95 folly::SocketAddress address_;
97 void connectSuccess() noexcept override {}
98 void connectErr(const folly::AsyncSocketException& ex) noexcept override {
101 void getReadBuffer(void** bufReturn, size_t* lenReturn) override {
102 *bufReturn = readBuf_;
103 *lenReturn = readLen_;
105 void readDataAvailable(size_t len) noexcept override {
109 sock_->setReadCB(nullptr);
112 void readEOF() noexcept override {}
113 void readErr(const folly::AsyncSocketException& ex) noexcept override {
116 void writeSuccess() noexcept override {}
118 size_t /* bytesWritten */,
119 const folly::AsyncSocketException& ex) noexcept override {
123 int32_t readHelper(uint8_t* buf, size_t len, bool all) {
124 if (!sock_->good()) {
130 sock_->setReadCB(this);
131 while (!err_ && sock_->good() && readLen_ > 0) {
132 eventBase_.loopOnce();
137 sock_->setReadCB(nullptr);
138 if (err_.hasValue()) {
141 if (all && readLen_ > 0) {
142 throw folly::AsyncSocketException(
143 folly::AsyncSocketException::UNKNOWN, "eof");
145 return len - readLen_;