In AsyncSocket::handleErrMessages() inside the loop check if the callback was not...
[folly.git] / folly / io / async / AsyncSocket.cpp
1 /*
2  * Copyright 2017 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
17 #include <folly/io/async/AsyncSocket.h>
18
19 #include <folly/ExceptionWrapper.h>
20 #include <folly/Portability.h>
21 #include <folly/SocketAddress.h>
22 #include <folly/io/Cursor.h>
23 #include <folly/io/IOBuf.h>
24 #include <folly/io/IOBufQueue.h>
25 #include <folly/portability/Fcntl.h>
26 #include <folly/portability/Sockets.h>
27 #include <folly/portability/SysUio.h>
28 #include <folly/portability/Unistd.h>
29
30 #include <boost/preprocessor/control/if.hpp>
31 #include <errno.h>
32 #include <limits.h>
33 #include <sys/types.h>
34 #include <thread>
35
36 using std::string;
37 using std::unique_ptr;
38
39 namespace fsp = folly::portability::sockets;
40
41 namespace folly {
42
43 static constexpr bool msgErrQueueSupported =
44 #ifdef MSG_ERRQUEUE
45     true;
46 #else
47     false;
48 #endif // MSG_ERRQUEUE
49
50 // static members initializers
51 const AsyncSocket::OptionMap AsyncSocket::emptyOptionMap;
52
53 const AsyncSocketException socketClosedLocallyEx(
54     AsyncSocketException::END_OF_FILE, "socket closed locally");
55 const AsyncSocketException socketShutdownForWritesEx(
56     AsyncSocketException::END_OF_FILE, "socket shutdown for writes");
57
58 // TODO: It might help performance to provide a version of BytesWriteRequest that
59 // users could derive from, so we can avoid the extra allocation for each call
60 // to write()/writev().  We could templatize TFramedAsyncChannel just like the
61 // protocols are currently templatized for transports.
62 //
63 // We would need the version for external users where they provide the iovec
64 // storage space, and only our internal version would allocate it at the end of
65 // the WriteRequest.
66
67 /* The default WriteRequest implementation, used for write(), writev() and
68  * writeChain()
69  *
70  * A new BytesWriteRequest operation is allocated on the heap for all write
71  * operations that cannot be completed immediately.
72  */
73 class AsyncSocket::BytesWriteRequest : public AsyncSocket::WriteRequest {
74  public:
75   static BytesWriteRequest* newRequest(AsyncSocket* socket,
76                                        WriteCallback* callback,
77                                        const iovec* ops,
78                                        uint32_t opCount,
79                                        uint32_t partialWritten,
80                                        uint32_t bytesWritten,
81                                        unique_ptr<IOBuf>&& ioBuf,
82                                        WriteFlags flags) {
83     assert(opCount > 0);
84     // Since we put a variable size iovec array at the end
85     // of each BytesWriteRequest, we have to manually allocate the memory.
86     void* buf = malloc(sizeof(BytesWriteRequest) +
87                        (opCount * sizeof(struct iovec)));
88     if (buf == nullptr) {
89       throw std::bad_alloc();
90     }
91
92     return new(buf) BytesWriteRequest(socket, callback, ops, opCount,
93                                       partialWritten, bytesWritten,
94                                       std::move(ioBuf), flags);
95   }
96
97   void destroy() override {
98     this->~BytesWriteRequest();
99     free(this);
100   }
101
102   WriteResult performWrite() override {
103     WriteFlags writeFlags = flags_;
104     if (getNext() != nullptr) {
105       writeFlags |= WriteFlags::CORK;
106     }
107     auto writeResult = socket_->performWrite(
108         getOps(), getOpCount(), writeFlags, &opsWritten_, &partialBytes_);
109     bytesWritten_ = writeResult.writeReturn > 0 ? writeResult.writeReturn : 0;
110     return writeResult;
111   }
112
113   bool isComplete() override {
114     return opsWritten_ == getOpCount();
115   }
116
117   void consume() override {
118     // Advance opIndex_ forward by opsWritten_
119     opIndex_ += opsWritten_;
120     assert(opIndex_ < opCount_);
121
122     // If we've finished writing any IOBufs, release them
123     if (ioBuf_) {
124       for (uint32_t i = opsWritten_; i != 0; --i) {
125         assert(ioBuf_);
126         ioBuf_ = ioBuf_->pop();
127       }
128     }
129
130     // Move partialBytes_ forward into the current iovec buffer
131     struct iovec* currentOp = writeOps_ + opIndex_;
132     assert((partialBytes_ < currentOp->iov_len) || (currentOp->iov_len == 0));
133     currentOp->iov_base =
134       reinterpret_cast<uint8_t*>(currentOp->iov_base) + partialBytes_;
135     currentOp->iov_len -= partialBytes_;
136
137     // Increment the totalBytesWritten_ count by bytesWritten_;
138     assert(bytesWritten_ >= 0);
139     totalBytesWritten_ += uint32_t(bytesWritten_);
140   }
141
142  private:
143   BytesWriteRequest(AsyncSocket* socket,
144                     WriteCallback* callback,
145                     const struct iovec* ops,
146                     uint32_t opCount,
147                     uint32_t partialBytes,
148                     uint32_t bytesWritten,
149                     unique_ptr<IOBuf>&& ioBuf,
150                     WriteFlags flags)
151     : AsyncSocket::WriteRequest(socket, callback)
152     , opCount_(opCount)
153     , opIndex_(0)
154     , flags_(flags)
155     , ioBuf_(std::move(ioBuf))
156     , opsWritten_(0)
157     , partialBytes_(partialBytes)
158     , bytesWritten_(bytesWritten) {
159     memcpy(writeOps_, ops, sizeof(*ops) * opCount_);
160   }
161
162   // private destructor, to ensure callers use destroy()
163   ~BytesWriteRequest() override = default;
164
165   const struct iovec* getOps() const {
166     assert(opCount_ > opIndex_);
167     return writeOps_ + opIndex_;
168   }
169
170   uint32_t getOpCount() const {
171     assert(opCount_ > opIndex_);
172     return opCount_ - opIndex_;
173   }
174
175   uint32_t opCount_;            ///< number of entries in writeOps_
176   uint32_t opIndex_;            ///< current index into writeOps_
177   WriteFlags flags_;            ///< set for WriteFlags
178   unique_ptr<IOBuf> ioBuf_;     ///< underlying IOBuf, or nullptr if N/A
179
180   // for consume(), how much we wrote on the last write
181   uint32_t opsWritten_;         ///< complete ops written
182   uint32_t partialBytes_;       ///< partial bytes of incomplete op written
183   ssize_t bytesWritten_;        ///< bytes written altogether
184
185   struct iovec writeOps_[];     ///< write operation(s) list
186 };
187
188 int AsyncSocket::SendMsgParamsCallback::getDefaultFlags(folly::WriteFlags flags)
189                                                                       noexcept {
190   int msg_flags = MSG_DONTWAIT;
191
192 #ifdef MSG_NOSIGNAL // Linux-only
193   msg_flags |= MSG_NOSIGNAL;
194 #ifdef MSG_MORE
195   if (isSet(flags, WriteFlags::CORK)) {
196     // MSG_MORE tells the kernel we have more data to send, so wait for us to
197     // give it the rest of the data rather than immediately sending a partial
198     // frame, even when TCP_NODELAY is enabled.
199     msg_flags |= MSG_MORE;
200   }
201 #endif // MSG_MORE
202 #endif // MSG_NOSIGNAL
203   if (isSet(flags, WriteFlags::EOR)) {
204     // marks that this is the last byte of a record (response)
205     msg_flags |= MSG_EOR;
206   }
207
208   return msg_flags;
209 }
210
211 namespace {
212 static AsyncSocket::SendMsgParamsCallback defaultSendMsgParamsCallback;
213 }
214
215 AsyncSocket::AsyncSocket()
216     : eventBase_(nullptr),
217       writeTimeout_(this, nullptr),
218       ioHandler_(this, nullptr),
219       immediateReadHandler_(this) {
220   VLOG(5) << "new AsyncSocket()";
221   init();
222 }
223
224 AsyncSocket::AsyncSocket(EventBase* evb)
225     : eventBase_(evb),
226       writeTimeout_(this, evb),
227       ioHandler_(this, evb),
228       immediateReadHandler_(this) {
229   VLOG(5) << "new AsyncSocket(" << this << ", evb=" << evb << ")";
230   init();
231 }
232
233 AsyncSocket::AsyncSocket(EventBase* evb,
234                            const folly::SocketAddress& address,
235                            uint32_t connectTimeout)
236   : AsyncSocket(evb) {
237   connect(nullptr, address, connectTimeout);
238 }
239
240 AsyncSocket::AsyncSocket(EventBase* evb,
241                            const std::string& ip,
242                            uint16_t port,
243                            uint32_t connectTimeout)
244   : AsyncSocket(evb) {
245   connect(nullptr, ip, port, connectTimeout);
246 }
247
248 AsyncSocket::AsyncSocket(EventBase* evb, int fd)
249     : eventBase_(evb),
250       writeTimeout_(this, evb),
251       ioHandler_(this, evb, fd),
252       immediateReadHandler_(this) {
253   VLOG(5) << "new AsyncSocket(" << this << ", evb=" << evb << ", fd="
254           << fd << ")";
255   init();
256   fd_ = fd;
257   setCloseOnExec();
258   state_ = StateEnum::ESTABLISHED;
259 }
260
261 AsyncSocket::AsyncSocket(AsyncSocket::UniquePtr oldAsyncSocket)
262     : AsyncSocket(oldAsyncSocket->getEventBase(), oldAsyncSocket->detachFd()) {
263   preReceivedData_ = std::move(oldAsyncSocket->preReceivedData_);
264 }
265
266 // init() method, since constructor forwarding isn't supported in most
267 // compilers yet.
268 void AsyncSocket::init() {
269   assert(eventBase_ == nullptr || eventBase_->isInEventBaseThread());
270   shutdownFlags_ = 0;
271   state_ = StateEnum::UNINIT;
272   eventFlags_ = EventHandler::NONE;
273   fd_ = -1;
274   sendTimeout_ = 0;
275   maxReadsPerEvent_ = 16;
276   connectCallback_ = nullptr;
277   errMessageCallback_ = nullptr;
278   readCallback_ = nullptr;
279   writeReqHead_ = nullptr;
280   writeReqTail_ = nullptr;
281   shutdownSocketSet_ = nullptr;
282   appBytesWritten_ = 0;
283   appBytesReceived_ = 0;
284   sendMsgParamCallback_ = &defaultSendMsgParamsCallback;
285 }
286
287 AsyncSocket::~AsyncSocket() {
288   VLOG(7) << "actual destruction of AsyncSocket(this=" << this
289           << ", evb=" << eventBase_ << ", fd=" << fd_
290           << ", state=" << state_ << ")";
291 }
292
293 void AsyncSocket::destroy() {
294   VLOG(5) << "AsyncSocket::destroy(this=" << this << ", evb=" << eventBase_
295           << ", fd=" << fd_ << ", state=" << state_;
296   // When destroy is called, close the socket immediately
297   closeNow();
298
299   // Then call DelayedDestruction::destroy() to take care of
300   // whether or not we need immediate or delayed destruction
301   DelayedDestruction::destroy();
302 }
303
304 int AsyncSocket::detachFd() {
305   VLOG(6) << "AsyncSocket::detachFd(this=" << this << ", fd=" << fd_
306           << ", evb=" << eventBase_ << ", state=" << state_
307           << ", events=" << std::hex << eventFlags_ << ")";
308   // Extract the fd, and set fd_ to -1 first, so closeNow() won't
309   // actually close the descriptor.
310   if (shutdownSocketSet_) {
311     shutdownSocketSet_->remove(fd_);
312   }
313   int fd = fd_;
314   fd_ = -1;
315   // Call closeNow() to invoke all pending callbacks with an error.
316   closeNow();
317   // Update the EventHandler to stop using this fd.
318   // This can only be done after closeNow() unregisters the handler.
319   ioHandler_.changeHandlerFD(-1);
320   return fd;
321 }
322
323 const folly::SocketAddress& AsyncSocket::anyAddress() {
324   static const folly::SocketAddress anyAddress =
325     folly::SocketAddress("0.0.0.0", 0);
326   return anyAddress;
327 }
328
329 void AsyncSocket::setShutdownSocketSet(ShutdownSocketSet* newSS) {
330   if (shutdownSocketSet_ == newSS) {
331     return;
332   }
333   if (shutdownSocketSet_ && fd_ != -1) {
334     shutdownSocketSet_->remove(fd_);
335   }
336   shutdownSocketSet_ = newSS;
337   if (shutdownSocketSet_ && fd_ != -1) {
338     shutdownSocketSet_->add(fd_);
339   }
340 }
341
342 void AsyncSocket::setCloseOnExec() {
343   int rv = fcntl(fd_, F_SETFD, FD_CLOEXEC);
344   if (rv != 0) {
345     auto errnoCopy = errno;
346     throw AsyncSocketException(
347         AsyncSocketException::INTERNAL_ERROR,
348         withAddr("failed to set close-on-exec flag"),
349         errnoCopy);
350   }
351 }
352
353 void AsyncSocket::connect(ConnectCallback* callback,
354                            const folly::SocketAddress& address,
355                            int timeout,
356                            const OptionMap &options,
357                            const folly::SocketAddress& bindAddr) noexcept {
358   DestructorGuard dg(this);
359   assert(eventBase_->isInEventBaseThread());
360
361   addr_ = address;
362
363   // Make sure we're in the uninitialized state
364   if (state_ != StateEnum::UNINIT) {
365     return invalidState(callback);
366   }
367
368   connectTimeout_ = std::chrono::milliseconds(timeout);
369   connectStartTime_ = std::chrono::steady_clock::now();
370   // Make connect end time at least >= connectStartTime.
371   connectEndTime_ = connectStartTime_;
372
373   assert(fd_ == -1);
374   state_ = StateEnum::CONNECTING;
375   connectCallback_ = callback;
376
377   sockaddr_storage addrStorage;
378   sockaddr* saddr = reinterpret_cast<sockaddr*>(&addrStorage);
379
380   try {
381     // Create the socket
382     // Technically the first parameter should actually be a protocol family
383     // constant (PF_xxx) rather than an address family (AF_xxx), but the
384     // distinction is mainly just historical.  In pretty much all
385     // implementations the PF_foo and AF_foo constants are identical.
386     fd_ = fsp::socket(address.getFamily(), SOCK_STREAM, 0);
387     if (fd_ < 0) {
388       auto errnoCopy = errno;
389       throw AsyncSocketException(
390           AsyncSocketException::INTERNAL_ERROR,
391           withAddr("failed to create socket"),
392           errnoCopy);
393     }
394     if (shutdownSocketSet_) {
395       shutdownSocketSet_->add(fd_);
396     }
397     ioHandler_.changeHandlerFD(fd_);
398
399     setCloseOnExec();
400
401     // Put the socket in non-blocking mode
402     int flags = fcntl(fd_, F_GETFL, 0);
403     if (flags == -1) {
404       auto errnoCopy = errno;
405       throw AsyncSocketException(
406           AsyncSocketException::INTERNAL_ERROR,
407           withAddr("failed to get socket flags"),
408           errnoCopy);
409     }
410     int rv = fcntl(fd_, F_SETFL, flags | O_NONBLOCK);
411     if (rv == -1) {
412       auto errnoCopy = errno;
413       throw AsyncSocketException(
414           AsyncSocketException::INTERNAL_ERROR,
415           withAddr("failed to put socket in non-blocking mode"),
416           errnoCopy);
417     }
418
419 #if !defined(MSG_NOSIGNAL) && defined(F_SETNOSIGPIPE)
420     // iOS and OS X don't support MSG_NOSIGNAL; set F_SETNOSIGPIPE instead
421     rv = fcntl(fd_, F_SETNOSIGPIPE, 1);
422     if (rv == -1) {
423       auto errnoCopy = errno;
424       throw AsyncSocketException(
425           AsyncSocketException::INTERNAL_ERROR,
426           "failed to enable F_SETNOSIGPIPE on socket",
427           errnoCopy);
428     }
429 #endif
430
431     // By default, turn on TCP_NODELAY
432     // If setNoDelay() fails, we continue anyway; this isn't a fatal error.
433     // setNoDelay() will log an error message if it fails.
434     if (address.getFamily() != AF_UNIX) {
435       (void)setNoDelay(true);
436     }
437
438     VLOG(5) << "AsyncSocket::connect(this=" << this << ", evb=" << eventBase_
439             << ", fd=" << fd_ << ", host=" << address.describe().c_str();
440
441     // bind the socket
442     if (bindAddr != anyAddress()) {
443       int one = 1;
444       if (setsockopt(fd_, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one))) {
445         auto errnoCopy = errno;
446         doClose();
447         throw AsyncSocketException(
448             AsyncSocketException::NOT_OPEN,
449             "failed to setsockopt prior to bind on " + bindAddr.describe(),
450             errnoCopy);
451       }
452
453       bindAddr.getAddress(&addrStorage);
454
455       if (bind(fd_, saddr, bindAddr.getActualSize()) != 0) {
456         auto errnoCopy = errno;
457         doClose();
458         throw AsyncSocketException(
459             AsyncSocketException::NOT_OPEN,
460             "failed to bind to async socket: " + bindAddr.describe(),
461             errnoCopy);
462       }
463     }
464
465     // Apply the additional options if any.
466     for (const auto& opt: options) {
467       rv = opt.first.apply(fd_, opt.second);
468       if (rv != 0) {
469         auto errnoCopy = errno;
470         throw AsyncSocketException(
471             AsyncSocketException::INTERNAL_ERROR,
472             withAddr("failed to set socket option"),
473             errnoCopy);
474       }
475     }
476
477     // Perform the connect()
478     address.getAddress(&addrStorage);
479
480     if (tfoEnabled_) {
481       state_ = StateEnum::FAST_OPEN;
482       tfoAttempted_ = true;
483     } else {
484       if (socketConnect(saddr, addr_.getActualSize()) < 0) {
485         return;
486       }
487     }
488
489     // If we're still here the connect() succeeded immediately.
490     // Fall through to call the callback outside of this try...catch block
491   } catch (const AsyncSocketException& ex) {
492     return failConnect(__func__, ex);
493   } catch (const std::exception& ex) {
494     // shouldn't happen, but handle it just in case
495     VLOG(4) << "AsyncSocket::connect(this=" << this << ", fd=" << fd_
496                << "): unexpected " << typeid(ex).name() << " exception: "
497                << ex.what();
498     AsyncSocketException tex(AsyncSocketException::INTERNAL_ERROR,
499                             withAddr(string("unexpected exception: ") +
500                                      ex.what()));
501     return failConnect(__func__, tex);
502   }
503
504   // The connection succeeded immediately
505   // The read callback may not have been set yet, and no writes may be pending
506   // yet, so we don't have to register for any events at the moment.
507   VLOG(8) << "AsyncSocket::connect succeeded immediately; this=" << this;
508   assert(errMessageCallback_ == nullptr);
509   assert(readCallback_ == nullptr);
510   assert(writeReqHead_ == nullptr);
511   if (state_ != StateEnum::FAST_OPEN) {
512     state_ = StateEnum::ESTABLISHED;
513   }
514   invokeConnectSuccess();
515 }
516
517 int AsyncSocket::socketConnect(const struct sockaddr* saddr, socklen_t len) {
518 #if __linux__
519   if (noTransparentTls_) {
520     // Ignore return value, errors are ok
521     setsockopt(fd_, SOL_SOCKET, SO_NO_TRANSPARENT_TLS, nullptr, 0);
522   }
523 #endif
524   int rv = fsp::connect(fd_, saddr, len);
525   if (rv < 0) {
526     auto errnoCopy = errno;
527     if (errnoCopy == EINPROGRESS) {
528       scheduleConnectTimeout();
529       registerForConnectEvents();
530     } else {
531       throw AsyncSocketException(
532           AsyncSocketException::NOT_OPEN,
533           "connect failed (immediately)",
534           errnoCopy);
535     }
536   }
537   return rv;
538 }
539
540 void AsyncSocket::scheduleConnectTimeout() {
541   // Connection in progress.
542   auto timeout = connectTimeout_.count();
543   if (timeout > 0) {
544     // Start a timer in case the connection takes too long.
545     if (!writeTimeout_.scheduleTimeout(uint32_t(timeout))) {
546       throw AsyncSocketException(
547           AsyncSocketException::INTERNAL_ERROR,
548           withAddr("failed to schedule AsyncSocket connect timeout"));
549     }
550   }
551 }
552
553 void AsyncSocket::registerForConnectEvents() {
554   // Register for write events, so we'll
555   // be notified when the connection finishes/fails.
556   // Note that we don't register for a persistent event here.
557   assert(eventFlags_ == EventHandler::NONE);
558   eventFlags_ = EventHandler::WRITE;
559   if (!ioHandler_.registerHandler(eventFlags_)) {
560     throw AsyncSocketException(
561         AsyncSocketException::INTERNAL_ERROR,
562         withAddr("failed to register AsyncSocket connect handler"));
563   }
564 }
565
566 void AsyncSocket::connect(ConnectCallback* callback,
567                            const string& ip, uint16_t port,
568                            int timeout,
569                            const OptionMap &options) noexcept {
570   DestructorGuard dg(this);
571   try {
572     connectCallback_ = callback;
573     connect(callback, folly::SocketAddress(ip, port), timeout, options);
574   } catch (const std::exception& ex) {
575     AsyncSocketException tex(AsyncSocketException::INTERNAL_ERROR,
576                             ex.what());
577     return failConnect(__func__, tex);
578   }
579 }
580
581 void AsyncSocket::cancelConnect() {
582   connectCallback_ = nullptr;
583   if (state_ == StateEnum::CONNECTING || state_ == StateEnum::FAST_OPEN) {
584     closeNow();
585   }
586 }
587
588 void AsyncSocket::setSendTimeout(uint32_t milliseconds) {
589   sendTimeout_ = milliseconds;
590   assert(eventBase_ == nullptr || eventBase_->isInEventBaseThread());
591
592   // If we are currently pending on write requests, immediately update
593   // writeTimeout_ with the new value.
594   if ((eventFlags_ & EventHandler::WRITE) &&
595       (state_ != StateEnum::CONNECTING && state_ != StateEnum::FAST_OPEN)) {
596     assert(state_ == StateEnum::ESTABLISHED);
597     assert((shutdownFlags_ & SHUT_WRITE) == 0);
598     if (sendTimeout_ > 0) {
599       if (!writeTimeout_.scheduleTimeout(sendTimeout_)) {
600         AsyncSocketException ex(AsyncSocketException::INTERNAL_ERROR,
601             withAddr("failed to reschedule send timeout in setSendTimeout"));
602         return failWrite(__func__, ex);
603       }
604     } else {
605       writeTimeout_.cancelTimeout();
606     }
607   }
608 }
609
610 void AsyncSocket::setErrMessageCB(ErrMessageCallback* callback) {
611   VLOG(6) << "AsyncSocket::setErrMessageCB() this=" << this
612           << ", fd=" << fd_ << ", callback=" << callback
613           << ", state=" << state_;
614
615   // Short circuit if callback is the same as the existing errMessageCallback_.
616   if (callback == errMessageCallback_) {
617     return;
618   }
619
620   if (!msgErrQueueSupported) {
621       // Per-socket error message queue is not supported on this platform.
622       return invalidState(callback);
623   }
624
625   DestructorGuard dg(this);
626   assert(eventBase_->isInEventBaseThread());
627
628   if (callback == nullptr) {
629     // We should be able to reset the callback regardless of the
630     // socket state. It's important to have a reliable callback
631     // cancellation mechanism.
632     errMessageCallback_ = callback;
633     return;
634   }
635
636   switch ((StateEnum)state_) {
637     case StateEnum::CONNECTING:
638     case StateEnum::FAST_OPEN:
639     case StateEnum::ESTABLISHED: {
640       errMessageCallback_ = callback;
641       return;
642     }
643     case StateEnum::CLOSED:
644     case StateEnum::ERROR:
645       // We should never reach here.  SHUT_READ should always be set
646       // if we are in STATE_CLOSED or STATE_ERROR.
647       assert(false);
648       return invalidState(callback);
649     case StateEnum::UNINIT:
650       // We do not allow setReadCallback() to be called before we start
651       // connecting.
652       return invalidState(callback);
653   }
654
655   // We don't put a default case in the switch statement, so that the compiler
656   // will warn us to update the switch statement if a new state is added.
657   return invalidState(callback);
658 }
659
660 AsyncSocket::ErrMessageCallback* AsyncSocket::getErrMessageCallback() const {
661   return errMessageCallback_;
662 }
663
664 void AsyncSocket::setSendMsgParamCB(SendMsgParamsCallback* callback) {
665   sendMsgParamCallback_ = callback;
666 }
667
668 AsyncSocket::SendMsgParamsCallback* AsyncSocket::getSendMsgParamsCB() const {
669   return sendMsgParamCallback_;
670 }
671
672 void AsyncSocket::setReadCB(ReadCallback *callback) {
673   VLOG(6) << "AsyncSocket::setReadCallback() this=" << this << ", fd=" << fd_
674           << ", callback=" << callback << ", state=" << state_;
675
676   // Short circuit if callback is the same as the existing readCallback_.
677   //
678   // Note that this is needed for proper functioning during some cleanup cases.
679   // During cleanup we allow setReadCallback(nullptr) to be called even if the
680   // read callback is already unset and we have been detached from an event
681   // base.  This check prevents us from asserting
682   // eventBase_->isInEventBaseThread() when eventBase_ is nullptr.
683   if (callback == readCallback_) {
684     return;
685   }
686
687   /* We are removing a read callback */
688   if (callback == nullptr &&
689       immediateReadHandler_.isLoopCallbackScheduled()) {
690     immediateReadHandler_.cancelLoopCallback();
691   }
692
693   if (shutdownFlags_ & SHUT_READ) {
694     // Reads have already been shut down on this socket.
695     //
696     // Allow setReadCallback(nullptr) to be called in this case, but don't
697     // allow a new callback to be set.
698     //
699     // For example, setReadCallback(nullptr) can happen after an error if we
700     // invoke some other error callback before invoking readError().  The other
701     // error callback that is invoked first may go ahead and clear the read
702     // callback before we get a chance to invoke readError().
703     if (callback != nullptr) {
704       return invalidState(callback);
705     }
706     assert((eventFlags_ & EventHandler::READ) == 0);
707     readCallback_ = nullptr;
708     return;
709   }
710
711   DestructorGuard dg(this);
712   assert(eventBase_->isInEventBaseThread());
713
714   switch ((StateEnum)state_) {
715     case StateEnum::CONNECTING:
716     case StateEnum::FAST_OPEN:
717       // For convenience, we allow the read callback to be set while we are
718       // still connecting.  We just store the callback for now.  Once the
719       // connection completes we'll register for read events.
720       readCallback_ = callback;
721       return;
722     case StateEnum::ESTABLISHED:
723     {
724       readCallback_ = callback;
725       uint16_t oldFlags = eventFlags_;
726       if (readCallback_) {
727         eventFlags_ |= EventHandler::READ;
728       } else {
729         eventFlags_ &= ~EventHandler::READ;
730       }
731
732       // Update our registration if our flags have changed
733       if (eventFlags_ != oldFlags) {
734         // We intentionally ignore the return value here.
735         // updateEventRegistration() will move us into the error state if it
736         // fails, and we don't need to do anything else here afterwards.
737         (void)updateEventRegistration();
738       }
739
740       if (readCallback_) {
741         checkForImmediateRead();
742       }
743       return;
744     }
745     case StateEnum::CLOSED:
746     case StateEnum::ERROR:
747       // We should never reach here.  SHUT_READ should always be set
748       // if we are in STATE_CLOSED or STATE_ERROR.
749       assert(false);
750       return invalidState(callback);
751     case StateEnum::UNINIT:
752       // We do not allow setReadCallback() to be called before we start
753       // connecting.
754       return invalidState(callback);
755   }
756
757   // We don't put a default case in the switch statement, so that the compiler
758   // will warn us to update the switch statement if a new state is added.
759   return invalidState(callback);
760 }
761
762 AsyncSocket::ReadCallback* AsyncSocket::getReadCallback() const {
763   return readCallback_;
764 }
765
766 void AsyncSocket::write(WriteCallback* callback,
767                          const void* buf, size_t bytes, WriteFlags flags) {
768   iovec op;
769   op.iov_base = const_cast<void*>(buf);
770   op.iov_len = bytes;
771   writeImpl(callback, &op, 1, unique_ptr<IOBuf>(), flags);
772 }
773
774 void AsyncSocket::writev(WriteCallback* callback,
775                           const iovec* vec,
776                           size_t count,
777                           WriteFlags flags) {
778   writeImpl(callback, vec, count, unique_ptr<IOBuf>(), flags);
779 }
780
781 void AsyncSocket::writeChain(WriteCallback* callback, unique_ptr<IOBuf>&& buf,
782                               WriteFlags flags) {
783   constexpr size_t kSmallSizeMax = 64;
784   size_t count = buf->countChainElements();
785   if (count <= kSmallSizeMax) {
786     // suppress "warning: variable length array 'vec' is used [-Wvla]"
787     FOLLY_PUSH_WARNING
788     FOLLY_GCC_DISABLE_WARNING("-Wvla")
789     iovec vec[BOOST_PP_IF(FOLLY_HAVE_VLA, count, kSmallSizeMax)];
790     FOLLY_POP_WARNING
791
792     writeChainImpl(callback, vec, count, std::move(buf), flags);
793   } else {
794     iovec* vec = new iovec[count];
795     writeChainImpl(callback, vec, count, std::move(buf), flags);
796     delete[] vec;
797   }
798 }
799
800 void AsyncSocket::writeChainImpl(WriteCallback* callback, iovec* vec,
801     size_t count, unique_ptr<IOBuf>&& buf, WriteFlags flags) {
802   size_t veclen = buf->fillIov(vec, count);
803   writeImpl(callback, vec, veclen, std::move(buf), flags);
804 }
805
806 void AsyncSocket::writeImpl(WriteCallback* callback, const iovec* vec,
807                              size_t count, unique_ptr<IOBuf>&& buf,
808                              WriteFlags flags) {
809   VLOG(6) << "AsyncSocket::writev() this=" << this << ", fd=" << fd_
810           << ", callback=" << callback << ", count=" << count
811           << ", state=" << state_;
812   DestructorGuard dg(this);
813   unique_ptr<IOBuf>ioBuf(std::move(buf));
814   assert(eventBase_->isInEventBaseThread());
815
816   if (shutdownFlags_ & (SHUT_WRITE | SHUT_WRITE_PENDING)) {
817     // No new writes may be performed after the write side of the socket has
818     // been shutdown.
819     //
820     // We could just call callback->writeError() here to fail just this write.
821     // However, fail hard and use invalidState() to fail all outstanding
822     // callbacks and move the socket into the error state.  There's most likely
823     // a bug in the caller's code, so we abort everything rather than trying to
824     // proceed as best we can.
825     return invalidState(callback);
826   }
827
828   uint32_t countWritten = 0;
829   uint32_t partialWritten = 0;
830   ssize_t bytesWritten = 0;
831   bool mustRegister = false;
832   if ((state_ == StateEnum::ESTABLISHED || state_ == StateEnum::FAST_OPEN) &&
833       !connecting()) {
834     if (writeReqHead_ == nullptr) {
835       // If we are established and there are no other writes pending,
836       // we can attempt to perform the write immediately.
837       assert(writeReqTail_ == nullptr);
838       assert((eventFlags_ & EventHandler::WRITE) == 0);
839
840       auto writeResult = performWrite(
841           vec, uint32_t(count), flags, &countWritten, &partialWritten);
842       bytesWritten = writeResult.writeReturn;
843       if (bytesWritten < 0) {
844         auto errnoCopy = errno;
845         if (writeResult.exception) {
846           return failWrite(__func__, callback, 0, *writeResult.exception);
847         }
848         AsyncSocketException ex(
849             AsyncSocketException::INTERNAL_ERROR,
850             withAddr("writev failed"),
851             errnoCopy);
852         return failWrite(__func__, callback, 0, ex);
853       } else if (countWritten == count) {
854         // We successfully wrote everything.
855         // Invoke the callback and return.
856         if (callback) {
857           callback->writeSuccess();
858         }
859         return;
860       } else { // continue writing the next writeReq
861         if (bufferCallback_) {
862           bufferCallback_->onEgressBuffered();
863         }
864       }
865       if (!connecting()) {
866         // Writes might put the socket back into connecting state
867         // if TFO is enabled, and using TFO fails.
868         // This means that write timeouts would not be active, however
869         // connect timeouts would affect this stage.
870         mustRegister = true;
871       }
872     }
873   } else if (!connecting()) {
874     // Invalid state for writing
875     return invalidState(callback);
876   }
877
878   // Create a new WriteRequest to add to the queue
879   WriteRequest* req;
880   try {
881     req = BytesWriteRequest::newRequest(
882         this,
883         callback,
884         vec + countWritten,
885         uint32_t(count - countWritten),
886         partialWritten,
887         uint32_t(bytesWritten),
888         std::move(ioBuf),
889         flags);
890   } catch (const std::exception& ex) {
891     // we mainly expect to catch std::bad_alloc here
892     AsyncSocketException tex(AsyncSocketException::INTERNAL_ERROR,
893         withAddr(string("failed to append new WriteRequest: ") + ex.what()));
894     return failWrite(__func__, callback, size_t(bytesWritten), tex);
895   }
896   req->consume();
897   if (writeReqTail_ == nullptr) {
898     assert(writeReqHead_ == nullptr);
899     writeReqHead_ = writeReqTail_ = req;
900   } else {
901     writeReqTail_->append(req);
902     writeReqTail_ = req;
903   }
904
905   // Register for write events if are established and not currently
906   // waiting on write events
907   if (mustRegister) {
908     assert(state_ == StateEnum::ESTABLISHED);
909     assert((eventFlags_ & EventHandler::WRITE) == 0);
910     if (!updateEventRegistration(EventHandler::WRITE, 0)) {
911       assert(state_ == StateEnum::ERROR);
912       return;
913     }
914     if (sendTimeout_ > 0) {
915       // Schedule a timeout to fire if the write takes too long.
916       if (!writeTimeout_.scheduleTimeout(sendTimeout_)) {
917         AsyncSocketException ex(AsyncSocketException::INTERNAL_ERROR,
918                                withAddr("failed to schedule send timeout"));
919         return failWrite(__func__, ex);
920       }
921     }
922   }
923 }
924
925 void AsyncSocket::writeRequest(WriteRequest* req) {
926   if (writeReqTail_ == nullptr) {
927     assert(writeReqHead_ == nullptr);
928     writeReqHead_ = writeReqTail_ = req;
929     req->start();
930   } else {
931     writeReqTail_->append(req);
932     writeReqTail_ = req;
933   }
934 }
935
936 void AsyncSocket::close() {
937   VLOG(5) << "AsyncSocket::close(): this=" << this << ", fd_=" << fd_
938           << ", state=" << state_ << ", shutdownFlags="
939           << std::hex << (int) shutdownFlags_;
940
941   // close() is only different from closeNow() when there are pending writes
942   // that need to drain before we can close.  In all other cases, just call
943   // closeNow().
944   //
945   // Note that writeReqHead_ can be non-nullptr even in STATE_CLOSED or
946   // STATE_ERROR if close() is invoked while a previous closeNow() or failure
947   // is still running.  (e.g., If there are multiple pending writes, and we
948   // call writeError() on the first one, it may call close().  In this case we
949   // will already be in STATE_CLOSED or STATE_ERROR, but the remaining pending
950   // writes will still be in the queue.)
951   //
952   // We only need to drain pending writes if we are still in STATE_CONNECTING
953   // or STATE_ESTABLISHED
954   if ((writeReqHead_ == nullptr) ||
955       !(state_ == StateEnum::CONNECTING ||
956       state_ == StateEnum::ESTABLISHED)) {
957     closeNow();
958     return;
959   }
960
961   // Declare a DestructorGuard to ensure that the AsyncSocket cannot be
962   // destroyed until close() returns.
963   DestructorGuard dg(this);
964   assert(eventBase_->isInEventBaseThread());
965
966   // Since there are write requests pending, we have to set the
967   // SHUT_WRITE_PENDING flag, and wait to perform the real close until the
968   // connect finishes and we finish writing these requests.
969   //
970   // Set SHUT_READ to indicate that reads are shut down, and set the
971   // SHUT_WRITE_PENDING flag to mark that we want to shutdown once the
972   // pending writes complete.
973   shutdownFlags_ |= (SHUT_READ | SHUT_WRITE_PENDING);
974
975   // If a read callback is set, invoke readEOF() immediately to inform it that
976   // the socket has been closed and no more data can be read.
977   if (readCallback_) {
978     // Disable reads if they are enabled
979     if (!updateEventRegistration(0, EventHandler::READ)) {
980       // We're now in the error state; callbacks have been cleaned up
981       assert(state_ == StateEnum::ERROR);
982       assert(readCallback_ == nullptr);
983     } else {
984       ReadCallback* callback = readCallback_;
985       readCallback_ = nullptr;
986       callback->readEOF();
987     }
988   }
989 }
990
991 void AsyncSocket::closeNow() {
992   VLOG(5) << "AsyncSocket::closeNow(): this=" << this << ", fd_=" << fd_
993           << ", state=" << state_ << ", shutdownFlags="
994           << std::hex << (int) shutdownFlags_;
995   DestructorGuard dg(this);
996   assert(eventBase_ == nullptr || eventBase_->isInEventBaseThread());
997
998   switch (state_) {
999     case StateEnum::ESTABLISHED:
1000     case StateEnum::CONNECTING:
1001     case StateEnum::FAST_OPEN: {
1002       shutdownFlags_ |= (SHUT_READ | SHUT_WRITE);
1003       state_ = StateEnum::CLOSED;
1004
1005       // If the write timeout was set, cancel it.
1006       writeTimeout_.cancelTimeout();
1007
1008       // If we are registered for I/O events, unregister.
1009       if (eventFlags_ != EventHandler::NONE) {
1010         eventFlags_ = EventHandler::NONE;
1011         if (!updateEventRegistration()) {
1012           // We will have been moved into the error state.
1013           assert(state_ == StateEnum::ERROR);
1014           return;
1015         }
1016       }
1017
1018       if (immediateReadHandler_.isLoopCallbackScheduled()) {
1019         immediateReadHandler_.cancelLoopCallback();
1020       }
1021
1022       if (fd_ >= 0) {
1023         ioHandler_.changeHandlerFD(-1);
1024         doClose();
1025       }
1026
1027       invokeConnectErr(socketClosedLocallyEx);
1028
1029       failAllWrites(socketClosedLocallyEx);
1030
1031       if (readCallback_) {
1032         ReadCallback* callback = readCallback_;
1033         readCallback_ = nullptr;
1034         callback->readEOF();
1035       }
1036       return;
1037     }
1038     case StateEnum::CLOSED:
1039       // Do nothing.  It's possible that we are being called recursively
1040       // from inside a callback that we invoked inside another call to close()
1041       // that is still running.
1042       return;
1043     case StateEnum::ERROR:
1044       // Do nothing.  The error handling code has performed (or is performing)
1045       // cleanup.
1046       return;
1047     case StateEnum::UNINIT:
1048       assert(eventFlags_ == EventHandler::NONE);
1049       assert(connectCallback_ == nullptr);
1050       assert(readCallback_ == nullptr);
1051       assert(writeReqHead_ == nullptr);
1052       shutdownFlags_ |= (SHUT_READ | SHUT_WRITE);
1053       state_ = StateEnum::CLOSED;
1054       return;
1055   }
1056
1057   LOG(DFATAL) << "AsyncSocket::closeNow() (this=" << this << ", fd=" << fd_
1058               << ") called in unknown state " << state_;
1059 }
1060
1061 void AsyncSocket::closeWithReset() {
1062   // Enable SO_LINGER, with the linger timeout set to 0.
1063   // This will trigger a TCP reset when we close the socket.
1064   if (fd_ >= 0) {
1065     struct linger optLinger = {1, 0};
1066     if (setSockOpt(SOL_SOCKET, SO_LINGER, &optLinger) != 0) {
1067       VLOG(2) << "AsyncSocket::closeWithReset(): error setting SO_LINGER "
1068               << "on " << fd_ << ": errno=" << errno;
1069     }
1070   }
1071
1072   // Then let closeNow() take care of the rest
1073   closeNow();
1074 }
1075
1076 void AsyncSocket::shutdownWrite() {
1077   VLOG(5) << "AsyncSocket::shutdownWrite(): this=" << this << ", fd=" << fd_
1078           << ", state=" << state_ << ", shutdownFlags="
1079           << std::hex << (int) shutdownFlags_;
1080
1081   // If there are no pending writes, shutdownWrite() is identical to
1082   // shutdownWriteNow().
1083   if (writeReqHead_ == nullptr) {
1084     shutdownWriteNow();
1085     return;
1086   }
1087
1088   assert(eventBase_->isInEventBaseThread());
1089
1090   // There are pending writes.  Set SHUT_WRITE_PENDING so that the actual
1091   // shutdown will be performed once all writes complete.
1092   shutdownFlags_ |= SHUT_WRITE_PENDING;
1093 }
1094
1095 void AsyncSocket::shutdownWriteNow() {
1096   VLOG(5) << "AsyncSocket::shutdownWriteNow(): this=" << this
1097           << ", fd=" << fd_ << ", state=" << state_
1098           << ", shutdownFlags=" << std::hex << (int) shutdownFlags_;
1099
1100   if (shutdownFlags_ & SHUT_WRITE) {
1101     // Writes are already shutdown; nothing else to do.
1102     return;
1103   }
1104
1105   // If SHUT_READ is already set, just call closeNow() to completely
1106   // close the socket.  This can happen if close() was called with writes
1107   // pending, and then shutdownWriteNow() is called before all pending writes
1108   // complete.
1109   if (shutdownFlags_ & SHUT_READ) {
1110     closeNow();
1111     return;
1112   }
1113
1114   DestructorGuard dg(this);
1115   assert(eventBase_ == nullptr || eventBase_->isInEventBaseThread());
1116
1117   switch (static_cast<StateEnum>(state_)) {
1118     case StateEnum::ESTABLISHED:
1119     {
1120       shutdownFlags_ |= SHUT_WRITE;
1121
1122       // If the write timeout was set, cancel it.
1123       writeTimeout_.cancelTimeout();
1124
1125       // If we are registered for write events, unregister.
1126       if (!updateEventRegistration(0, EventHandler::WRITE)) {
1127         // We will have been moved into the error state.
1128         assert(state_ == StateEnum::ERROR);
1129         return;
1130       }
1131
1132       // Shutdown writes on the file descriptor
1133       shutdown(fd_, SHUT_WR);
1134
1135       // Immediately fail all write requests
1136       failAllWrites(socketShutdownForWritesEx);
1137       return;
1138     }
1139     case StateEnum::CONNECTING:
1140     {
1141       // Set the SHUT_WRITE_PENDING flag.
1142       // When the connection completes, it will check this flag,
1143       // shutdown the write half of the socket, and then set SHUT_WRITE.
1144       shutdownFlags_ |= SHUT_WRITE_PENDING;
1145
1146       // Immediately fail all write requests
1147       failAllWrites(socketShutdownForWritesEx);
1148       return;
1149     }
1150     case StateEnum::UNINIT:
1151       // Callers normally shouldn't call shutdownWriteNow() before the socket
1152       // even starts connecting.  Nonetheless, go ahead and set
1153       // SHUT_WRITE_PENDING.  Once the socket eventually connects it will
1154       // immediately shut down the write side of the socket.
1155       shutdownFlags_ |= SHUT_WRITE_PENDING;
1156       return;
1157     case StateEnum::FAST_OPEN:
1158       // In fast open state we haven't call connected yet, and if we shutdown
1159       // the writes, we will never try to call connect, so shut everything down
1160       shutdownFlags_ |= SHUT_WRITE;
1161       // Immediately fail all write requests
1162       failAllWrites(socketShutdownForWritesEx);
1163       return;
1164     case StateEnum::CLOSED:
1165     case StateEnum::ERROR:
1166       // We should never get here.  SHUT_WRITE should always be set
1167       // in STATE_CLOSED and STATE_ERROR.
1168       VLOG(4) << "AsyncSocket::shutdownWriteNow() (this=" << this
1169                  << ", fd=" << fd_ << ") in unexpected state " << state_
1170                  << " with SHUT_WRITE not set ("
1171                  << std::hex << (int) shutdownFlags_ << ")";
1172       assert(false);
1173       return;
1174   }
1175
1176   LOG(DFATAL) << "AsyncSocket::shutdownWriteNow() (this=" << this << ", fd="
1177               << fd_ << ") called in unknown state " << state_;
1178 }
1179
1180 bool AsyncSocket::readable() const {
1181   if (fd_ == -1) {
1182     return false;
1183   }
1184   struct pollfd fds[1];
1185   fds[0].fd = fd_;
1186   fds[0].events = POLLIN;
1187   fds[0].revents = 0;
1188   int rc = poll(fds, 1, 0);
1189   return rc == 1;
1190 }
1191
1192 bool AsyncSocket::writable() const {
1193   if (fd_ == -1) {
1194     return false;
1195   }
1196   struct pollfd fds[1];
1197   fds[0].fd = fd_;
1198   fds[0].events = POLLOUT;
1199   fds[0].revents = 0;
1200   int rc = poll(fds, 1, 0);
1201   return rc == 1;
1202 }
1203
1204 bool AsyncSocket::isPending() const {
1205   return ioHandler_.isPending();
1206 }
1207
1208 bool AsyncSocket::hangup() const {
1209   if (fd_ == -1) {
1210     // sanity check, no one should ask for hangup if we are not connected.
1211     assert(false);
1212     return false;
1213   }
1214 #ifdef POLLRDHUP // Linux-only
1215   struct pollfd fds[1];
1216   fds[0].fd = fd_;
1217   fds[0].events = POLLRDHUP|POLLHUP;
1218   fds[0].revents = 0;
1219   poll(fds, 1, 0);
1220   return (fds[0].revents & (POLLRDHUP|POLLHUP)) != 0;
1221 #else
1222   return false;
1223 #endif
1224 }
1225
1226 bool AsyncSocket::good() const {
1227   return (
1228       (state_ == StateEnum::CONNECTING || state_ == StateEnum::FAST_OPEN ||
1229        state_ == StateEnum::ESTABLISHED) &&
1230       (shutdownFlags_ == 0) && (eventBase_ != nullptr));
1231 }
1232
1233 bool AsyncSocket::error() const {
1234   return (state_ == StateEnum::ERROR);
1235 }
1236
1237 void AsyncSocket::attachEventBase(EventBase* eventBase) {
1238   VLOG(5) << "AsyncSocket::attachEventBase(this=" << this << ", fd=" << fd_
1239           << ", old evb=" << eventBase_ << ", new evb=" << eventBase
1240           << ", state=" << state_ << ", events="
1241           << std::hex << eventFlags_ << ")";
1242   assert(eventBase_ == nullptr);
1243   assert(eventBase->isInEventBaseThread());
1244
1245   eventBase_ = eventBase;
1246   ioHandler_.attachEventBase(eventBase);
1247   writeTimeout_.attachEventBase(eventBase);
1248   if (evbChangeCb_) {
1249     evbChangeCb_->evbAttached(this);
1250   }
1251 }
1252
1253 void AsyncSocket::detachEventBase() {
1254   VLOG(5) << "AsyncSocket::detachEventBase(this=" << this << ", fd=" << fd_
1255           << ", old evb=" << eventBase_ << ", state=" << state_
1256           << ", events=" << std::hex << eventFlags_ << ")";
1257   assert(eventBase_ != nullptr);
1258   assert(eventBase_->isInEventBaseThread());
1259
1260   eventBase_ = nullptr;
1261   ioHandler_.detachEventBase();
1262   writeTimeout_.detachEventBase();
1263   if (evbChangeCb_) {
1264     evbChangeCb_->evbDetached(this);
1265   }
1266 }
1267
1268 bool AsyncSocket::isDetachable() const {
1269   DCHECK(eventBase_ != nullptr);
1270   DCHECK(eventBase_->isInEventBaseThread());
1271
1272   return !ioHandler_.isHandlerRegistered() && !writeTimeout_.isScheduled();
1273 }
1274
1275 void AsyncSocket::getLocalAddress(folly::SocketAddress* address) const {
1276   if (!localAddr_.isInitialized()) {
1277     localAddr_.setFromLocalAddress(fd_);
1278   }
1279   *address = localAddr_;
1280 }
1281
1282 void AsyncSocket::getPeerAddress(folly::SocketAddress* address) const {
1283   if (!addr_.isInitialized()) {
1284     addr_.setFromPeerAddress(fd_);
1285   }
1286   *address = addr_;
1287 }
1288
1289 bool AsyncSocket::getTFOSucceded() const {
1290   return detail::tfo_succeeded(fd_);
1291 }
1292
1293 int AsyncSocket::setNoDelay(bool noDelay) {
1294   if (fd_ < 0) {
1295     VLOG(4) << "AsyncSocket::setNoDelay() called on non-open socket "
1296                << this << "(state=" << state_ << ")";
1297     return EINVAL;
1298
1299   }
1300
1301   int value = noDelay ? 1 : 0;
1302   if (setsockopt(fd_, IPPROTO_TCP, TCP_NODELAY, &value, sizeof(value)) != 0) {
1303     int errnoCopy = errno;
1304     VLOG(2) << "failed to update TCP_NODELAY option on AsyncSocket "
1305             << this << " (fd=" << fd_ << ", state=" << state_ << "): "
1306             << strerror(errnoCopy);
1307     return errnoCopy;
1308   }
1309
1310   return 0;
1311 }
1312
1313 int AsyncSocket::setCongestionFlavor(const std::string &cname) {
1314
1315   #ifndef TCP_CONGESTION
1316   #define TCP_CONGESTION  13
1317   #endif
1318
1319   if (fd_ < 0) {
1320     VLOG(4) << "AsyncSocket::setCongestionFlavor() called on non-open "
1321                << "socket " << this << "(state=" << state_ << ")";
1322     return EINVAL;
1323
1324   }
1325
1326   if (setsockopt(
1327           fd_,
1328           IPPROTO_TCP,
1329           TCP_CONGESTION,
1330           cname.c_str(),
1331           socklen_t(cname.length() + 1)) != 0) {
1332     int errnoCopy = errno;
1333     VLOG(2) << "failed to update TCP_CONGESTION option on AsyncSocket "
1334             << this << "(fd=" << fd_ << ", state=" << state_ << "): "
1335             << strerror(errnoCopy);
1336     return errnoCopy;
1337   }
1338
1339   return 0;
1340 }
1341
1342 int AsyncSocket::setQuickAck(bool quickack) {
1343   (void)quickack;
1344   if (fd_ < 0) {
1345     VLOG(4) << "AsyncSocket::setQuickAck() called on non-open socket "
1346                << this << "(state=" << state_ << ")";
1347     return EINVAL;
1348
1349   }
1350
1351 #ifdef TCP_QUICKACK // Linux-only
1352   int value = quickack ? 1 : 0;
1353   if (setsockopt(fd_, IPPROTO_TCP, TCP_QUICKACK, &value, sizeof(value)) != 0) {
1354     int errnoCopy = errno;
1355     VLOG(2) << "failed to update TCP_QUICKACK option on AsyncSocket"
1356             << this << "(fd=" << fd_ << ", state=" << state_ << "): "
1357             << strerror(errnoCopy);
1358     return errnoCopy;
1359   }
1360
1361   return 0;
1362 #else
1363   return ENOSYS;
1364 #endif
1365 }
1366
1367 int AsyncSocket::setSendBufSize(size_t bufsize) {
1368   if (fd_ < 0) {
1369     VLOG(4) << "AsyncSocket::setSendBufSize() called on non-open socket "
1370                << this << "(state=" << state_ << ")";
1371     return EINVAL;
1372   }
1373
1374   if (setsockopt(fd_, SOL_SOCKET, SO_SNDBUF, &bufsize, sizeof(bufsize)) !=0) {
1375     int errnoCopy = errno;
1376     VLOG(2) << "failed to update SO_SNDBUF option on AsyncSocket"
1377             << this << "(fd=" << fd_ << ", state=" << state_ << "): "
1378             << strerror(errnoCopy);
1379     return errnoCopy;
1380   }
1381
1382   return 0;
1383 }
1384
1385 int AsyncSocket::setRecvBufSize(size_t bufsize) {
1386   if (fd_ < 0) {
1387     VLOG(4) << "AsyncSocket::setRecvBufSize() called on non-open socket "
1388                << this << "(state=" << state_ << ")";
1389     return EINVAL;
1390   }
1391
1392   if (setsockopt(fd_, SOL_SOCKET, SO_RCVBUF, &bufsize, sizeof(bufsize)) !=0) {
1393     int errnoCopy = errno;
1394     VLOG(2) << "failed to update SO_RCVBUF option on AsyncSocket"
1395             << this << "(fd=" << fd_ << ", state=" << state_ << "): "
1396             << strerror(errnoCopy);
1397     return errnoCopy;
1398   }
1399
1400   return 0;
1401 }
1402
1403 int AsyncSocket::setTCPProfile(int profd) {
1404   if (fd_ < 0) {
1405     VLOG(4) << "AsyncSocket::setTCPProfile() called on non-open socket "
1406                << this << "(state=" << state_ << ")";
1407     return EINVAL;
1408   }
1409
1410   if (setsockopt(fd_, SOL_SOCKET, SO_SET_NAMESPACE, &profd, sizeof(int)) !=0) {
1411     int errnoCopy = errno;
1412     VLOG(2) << "failed to set socket namespace option on AsyncSocket"
1413             << this << "(fd=" << fd_ << ", state=" << state_ << "): "
1414             << strerror(errnoCopy);
1415     return errnoCopy;
1416   }
1417
1418   return 0;
1419 }
1420
1421 void AsyncSocket::ioReady(uint16_t events) noexcept {
1422   VLOG(7) << "AsyncSocket::ioRead() this=" << this << ", fd=" << fd_
1423           << ", events=" << std::hex << events << ", state=" << state_;
1424   DestructorGuard dg(this);
1425   assert(events & EventHandler::READ_WRITE);
1426   assert(eventBase_->isInEventBaseThread());
1427
1428   uint16_t relevantEvents = uint16_t(events & EventHandler::READ_WRITE);
1429   EventBase* originalEventBase = eventBase_;
1430   // If we got there it means that either EventHandler::READ or
1431   // EventHandler::WRITE is set. Any of these flags can
1432   // indicate that there are messages available in the socket
1433   // error message queue.
1434   handleErrMessages();
1435
1436   // Return now if handleErrMessages() detached us from our EventBase
1437   if (eventBase_ != originalEventBase) {
1438     return;
1439   }
1440
1441   if (relevantEvents == EventHandler::READ) {
1442     handleRead();
1443   } else if (relevantEvents == EventHandler::WRITE) {
1444     handleWrite();
1445   } else if (relevantEvents == EventHandler::READ_WRITE) {
1446     // If both read and write events are ready, process writes first.
1447     handleWrite();
1448
1449     // Return now if handleWrite() detached us from our EventBase
1450     if (eventBase_ != originalEventBase) {
1451       return;
1452     }
1453
1454     // Only call handleRead() if a read callback is still installed.
1455     // (It's possible that the read callback was uninstalled during
1456     // handleWrite().)
1457     if (readCallback_) {
1458       handleRead();
1459     }
1460   } else {
1461     VLOG(4) << "AsyncSocket::ioRead() called with unexpected events "
1462                << std::hex << events << "(this=" << this << ")";
1463     abort();
1464   }
1465 }
1466
1467 AsyncSocket::ReadResult
1468 AsyncSocket::performRead(void** buf, size_t* buflen, size_t* /* offset */) {
1469   VLOG(5) << "AsyncSocket::performRead() this=" << this << ", buf=" << *buf
1470           << ", buflen=" << *buflen;
1471
1472   if (preReceivedData_ && !preReceivedData_->empty()) {
1473     VLOG(5) << "AsyncSocket::performRead() this=" << this
1474             << ", reading pre-received data";
1475
1476     io::Cursor cursor(preReceivedData_.get());
1477     auto len = cursor.pullAtMost(*buf, *buflen);
1478
1479     IOBufQueue queue;
1480     queue.append(std::move(preReceivedData_));
1481     queue.trimStart(len);
1482     preReceivedData_ = queue.move();
1483
1484     appBytesReceived_ += len;
1485     return ReadResult(len);
1486   }
1487
1488   ssize_t bytes = recv(fd_, *buf, *buflen, MSG_DONTWAIT);
1489   if (bytes < 0) {
1490     if (errno == EAGAIN || errno == EWOULDBLOCK) {
1491       // No more data to read right now.
1492       return ReadResult(READ_BLOCKING);
1493     } else {
1494       return ReadResult(READ_ERROR);
1495     }
1496   } else {
1497     appBytesReceived_ += bytes;
1498     return ReadResult(bytes);
1499   }
1500 }
1501
1502 void AsyncSocket::prepareReadBuffer(void** buf, size_t* buflen) {
1503   // no matter what, buffer should be preapared for non-ssl socket
1504   CHECK(readCallback_);
1505   readCallback_->getReadBuffer(buf, buflen);
1506 }
1507
1508 void AsyncSocket::handleErrMessages() noexcept {
1509   // This method has non-empty implementation only for platforms
1510   // supporting per-socket error queues.
1511   VLOG(5) << "AsyncSocket::handleErrMessages() this=" << this << ", fd=" << fd_
1512           << ", state=" << state_;
1513   if (errMessageCallback_ == nullptr) {
1514     VLOG(7) << "AsyncSocket::handleErrMessages(): "
1515             << "no callback installed - exiting.";
1516     return;
1517   }
1518
1519 #ifdef MSG_ERRQUEUE
1520   uint8_t ctrl[1024];
1521   unsigned char data;
1522   struct msghdr msg;
1523   iovec entry;
1524
1525   entry.iov_base = &data;
1526   entry.iov_len = sizeof(data);
1527   msg.msg_iov = &entry;
1528   msg.msg_iovlen = 1;
1529   msg.msg_name = nullptr;
1530   msg.msg_namelen = 0;
1531   msg.msg_control = ctrl;
1532   msg.msg_controllen = sizeof(ctrl);
1533   msg.msg_flags = 0;
1534
1535   int ret;
1536   while (true) {
1537     ret = recvmsg(fd_, &msg, MSG_ERRQUEUE);
1538     VLOG(5) << "AsyncSocket::handleErrMessages(): recvmsg returned " << ret;
1539
1540     if (ret < 0) {
1541       if (errno != EAGAIN) {
1542         auto errnoCopy = errno;
1543         LOG(ERROR) << "::recvmsg exited with code " << ret
1544                    << ", errno: " << errnoCopy;
1545         AsyncSocketException ex(
1546           AsyncSocketException::INTERNAL_ERROR,
1547           withAddr("recvmsg() failed"),
1548           errnoCopy);
1549         failErrMessageRead(__func__, ex);
1550       }
1551       return;
1552     }
1553
1554     for (struct cmsghdr* cmsg = CMSG_FIRSTHDR(&msg);
1555          cmsg != nullptr &&
1556            cmsg->cmsg_len != 0 &&
1557            errMessageCallback_ != nullptr;
1558          cmsg = CMSG_NXTHDR(&msg, cmsg)) {
1559       errMessageCallback_->errMessage(*cmsg);
1560     }
1561   }
1562 #endif //MSG_ERRQUEUE
1563 }
1564
1565 void AsyncSocket::handleRead() noexcept {
1566   VLOG(5) << "AsyncSocket::handleRead() this=" << this << ", fd=" << fd_
1567           << ", state=" << state_;
1568   assert(state_ == StateEnum::ESTABLISHED);
1569   assert((shutdownFlags_ & SHUT_READ) == 0);
1570   assert(readCallback_ != nullptr);
1571   assert(eventFlags_ & EventHandler::READ);
1572
1573   // Loop until:
1574   // - a read attempt would block
1575   // - readCallback_ is uninstalled
1576   // - the number of loop iterations exceeds the optional maximum
1577   // - this AsyncSocket is moved to another EventBase
1578   //
1579   // When we invoke readDataAvailable() it may uninstall the readCallback_,
1580   // which is why need to check for it here.
1581   //
1582   // The last bullet point is slightly subtle.  readDataAvailable() may also
1583   // detach this socket from this EventBase.  However, before
1584   // readDataAvailable() returns another thread may pick it up, attach it to
1585   // a different EventBase, and install another readCallback_.  We need to
1586   // exit immediately after readDataAvailable() returns if the eventBase_ has
1587   // changed.  (The caller must perform some sort of locking to transfer the
1588   // AsyncSocket between threads properly.  This will be sufficient to ensure
1589   // that this thread sees the updated eventBase_ variable after
1590   // readDataAvailable() returns.)
1591   uint16_t numReads = 0;
1592   EventBase* originalEventBase = eventBase_;
1593   while (readCallback_ && eventBase_ == originalEventBase) {
1594     // Get the buffer to read into.
1595     void* buf = nullptr;
1596     size_t buflen = 0, offset = 0;
1597     try {
1598       prepareReadBuffer(&buf, &buflen);
1599       VLOG(5) << "prepareReadBuffer() buf=" << buf << ", buflen=" << buflen;
1600     } catch (const AsyncSocketException& ex) {
1601       return failRead(__func__, ex);
1602     } catch (const std::exception& ex) {
1603       AsyncSocketException tex(AsyncSocketException::BAD_ARGS,
1604                               string("ReadCallback::getReadBuffer() "
1605                                      "threw exception: ") +
1606                               ex.what());
1607       return failRead(__func__, tex);
1608     } catch (...) {
1609       AsyncSocketException ex(AsyncSocketException::BAD_ARGS,
1610                              "ReadCallback::getReadBuffer() threw "
1611                              "non-exception type");
1612       return failRead(__func__, ex);
1613     }
1614     if (!isBufferMovable_ && (buf == nullptr || buflen == 0)) {
1615       AsyncSocketException ex(AsyncSocketException::BAD_ARGS,
1616                              "ReadCallback::getReadBuffer() returned "
1617                              "empty buffer");
1618       return failRead(__func__, ex);
1619     }
1620
1621     // Perform the read
1622     auto readResult = performRead(&buf, &buflen, &offset);
1623     auto bytesRead = readResult.readReturn;
1624     VLOG(4) << "this=" << this << ", AsyncSocket::handleRead() got "
1625             << bytesRead << " bytes";
1626     if (bytesRead > 0) {
1627       if (!isBufferMovable_) {
1628         readCallback_->readDataAvailable(size_t(bytesRead));
1629       } else {
1630         CHECK(kOpenSslModeMoveBufferOwnership);
1631         VLOG(5) << "this=" << this << ", AsyncSocket::handleRead() got "
1632                 << "buf=" << buf << ", " << bytesRead << "/" << buflen
1633                 << ", offset=" << offset;
1634         auto readBuf = folly::IOBuf::takeOwnership(buf, buflen);
1635         readBuf->trimStart(offset);
1636         readBuf->trimEnd(buflen - offset - bytesRead);
1637         readCallback_->readBufferAvailable(std::move(readBuf));
1638       }
1639
1640       // Fall through and continue around the loop if the read
1641       // completely filled the available buffer.
1642       // Note that readCallback_ may have been uninstalled or changed inside
1643       // readDataAvailable().
1644       if (size_t(bytesRead) < buflen) {
1645         return;
1646       }
1647     } else if (bytesRead == READ_BLOCKING) {
1648         // No more data to read right now.
1649         return;
1650     } else if (bytesRead == READ_ERROR) {
1651       readErr_ = READ_ERROR;
1652       if (readResult.exception) {
1653         return failRead(__func__, *readResult.exception);
1654       }
1655       auto errnoCopy = errno;
1656       AsyncSocketException ex(
1657           AsyncSocketException::INTERNAL_ERROR,
1658           withAddr("recv() failed"),
1659           errnoCopy);
1660       return failRead(__func__, ex);
1661     } else {
1662       assert(bytesRead == READ_EOF);
1663       readErr_ = READ_EOF;
1664       // EOF
1665       shutdownFlags_ |= SHUT_READ;
1666       if (!updateEventRegistration(0, EventHandler::READ)) {
1667         // we've already been moved into STATE_ERROR
1668         assert(state_ == StateEnum::ERROR);
1669         assert(readCallback_ == nullptr);
1670         return;
1671       }
1672
1673       ReadCallback* callback = readCallback_;
1674       readCallback_ = nullptr;
1675       callback->readEOF();
1676       return;
1677     }
1678     if (maxReadsPerEvent_ && (++numReads >= maxReadsPerEvent_)) {
1679       if (readCallback_ != nullptr) {
1680         // We might still have data in the socket.
1681         // (e.g. see comment in AsyncSSLSocket::checkForImmediateRead)
1682         scheduleImmediateRead();
1683       }
1684       return;
1685     }
1686   }
1687 }
1688
1689 /**
1690  * This function attempts to write as much data as possible, until no more data
1691  * can be written.
1692  *
1693  * - If it sends all available data, it unregisters for write events, and stops
1694  *   the writeTimeout_.
1695  *
1696  * - If not all of the data can be sent immediately, it reschedules
1697  *   writeTimeout_ (if a non-zero timeout is set), and ensures the handler is
1698  *   registered for write events.
1699  */
1700 void AsyncSocket::handleWrite() noexcept {
1701   VLOG(5) << "AsyncSocket::handleWrite() this=" << this << ", fd=" << fd_
1702           << ", state=" << state_;
1703   DestructorGuard dg(this);
1704
1705   if (state_ == StateEnum::CONNECTING) {
1706     handleConnect();
1707     return;
1708   }
1709
1710   // Normal write
1711   assert(state_ == StateEnum::ESTABLISHED);
1712   assert((shutdownFlags_ & SHUT_WRITE) == 0);
1713   assert(writeReqHead_ != nullptr);
1714
1715   // Loop until we run out of write requests,
1716   // or until this socket is moved to another EventBase.
1717   // (See the comment in handleRead() explaining how this can happen.)
1718   EventBase* originalEventBase = eventBase_;
1719   while (writeReqHead_ != nullptr && eventBase_ == originalEventBase) {
1720     auto writeResult = writeReqHead_->performWrite();
1721     if (writeResult.writeReturn < 0) {
1722       if (writeResult.exception) {
1723         return failWrite(__func__, *writeResult.exception);
1724       }
1725       auto errnoCopy = errno;
1726       AsyncSocketException ex(
1727           AsyncSocketException::INTERNAL_ERROR,
1728           withAddr("writev() failed"),
1729           errnoCopy);
1730       return failWrite(__func__, ex);
1731     } else if (writeReqHead_->isComplete()) {
1732       // We finished this request
1733       WriteRequest* req = writeReqHead_;
1734       writeReqHead_ = req->getNext();
1735
1736       if (writeReqHead_ == nullptr) {
1737         writeReqTail_ = nullptr;
1738         // This is the last write request.
1739         // Unregister for write events and cancel the send timer
1740         // before we invoke the callback.  We have to update the state properly
1741         // before calling the callback, since it may want to detach us from
1742         // the EventBase.
1743         if (eventFlags_ & EventHandler::WRITE) {
1744           if (!updateEventRegistration(0, EventHandler::WRITE)) {
1745             assert(state_ == StateEnum::ERROR);
1746             return;
1747           }
1748           // Stop the send timeout
1749           writeTimeout_.cancelTimeout();
1750         }
1751         assert(!writeTimeout_.isScheduled());
1752
1753         // If SHUT_WRITE_PENDING is set, we should shutdown the socket after
1754         // we finish sending the last write request.
1755         //
1756         // We have to do this before invoking writeSuccess(), since
1757         // writeSuccess() may detach us from our EventBase.
1758         if (shutdownFlags_ & SHUT_WRITE_PENDING) {
1759           assert(connectCallback_ == nullptr);
1760           shutdownFlags_ |= SHUT_WRITE;
1761
1762           if (shutdownFlags_ & SHUT_READ) {
1763             // Reads have already been shutdown.  Fully close the socket and
1764             // move to STATE_CLOSED.
1765             //
1766             // Note: This code currently moves us to STATE_CLOSED even if
1767             // close() hasn't ever been called.  This can occur if we have
1768             // received EOF from the peer and shutdownWrite() has been called
1769             // locally.  Should we bother staying in STATE_ESTABLISHED in this
1770             // case, until close() is actually called?  I can't think of a
1771             // reason why we would need to do so.  No other operations besides
1772             // calling close() or destroying the socket can be performed at
1773             // this point.
1774             assert(readCallback_ == nullptr);
1775             state_ = StateEnum::CLOSED;
1776             if (fd_ >= 0) {
1777               ioHandler_.changeHandlerFD(-1);
1778               doClose();
1779             }
1780           } else {
1781             // Reads are still enabled, so we are only doing a half-shutdown
1782             shutdown(fd_, SHUT_WR);
1783           }
1784         }
1785       }
1786
1787       // Invoke the callback
1788       WriteCallback* callback = req->getCallback();
1789       req->destroy();
1790       if (callback) {
1791         callback->writeSuccess();
1792       }
1793       // We'll continue around the loop, trying to write another request
1794     } else {
1795       // Partial write.
1796       if (bufferCallback_) {
1797         bufferCallback_->onEgressBuffered();
1798       }
1799       writeReqHead_->consume();
1800       // Stop after a partial write; it's highly likely that a subsequent write
1801       // attempt will just return EAGAIN.
1802       //
1803       // Ensure that we are registered for write events.
1804       if ((eventFlags_ & EventHandler::WRITE) == 0) {
1805         if (!updateEventRegistration(EventHandler::WRITE, 0)) {
1806           assert(state_ == StateEnum::ERROR);
1807           return;
1808         }
1809       }
1810
1811       // Reschedule the send timeout, since we have made some write progress.
1812       if (sendTimeout_ > 0) {
1813         if (!writeTimeout_.scheduleTimeout(sendTimeout_)) {
1814           AsyncSocketException ex(AsyncSocketException::INTERNAL_ERROR,
1815               withAddr("failed to reschedule write timeout"));
1816           return failWrite(__func__, ex);
1817         }
1818       }
1819       return;
1820     }
1821   }
1822   if (!writeReqHead_ && bufferCallback_) {
1823     bufferCallback_->onEgressBufferCleared();
1824   }
1825 }
1826
1827 void AsyncSocket::checkForImmediateRead() noexcept {
1828   // We currently don't attempt to perform optimistic reads in AsyncSocket.
1829   // (However, note that some subclasses do override this method.)
1830   //
1831   // Simply calling handleRead() here would be bad, as this would call
1832   // readCallback_->getReadBuffer(), forcing the callback to allocate a read
1833   // buffer even though no data may be available.  This would waste lots of
1834   // memory, since the buffer will sit around unused until the socket actually
1835   // becomes readable.
1836   //
1837   // Checking if the socket is readable now also seems like it would probably
1838   // be a pessimism.  In most cases it probably wouldn't be readable, and we
1839   // would just waste an extra system call.  Even if it is readable, waiting to
1840   // find out from libevent on the next event loop doesn't seem that bad.
1841   //
1842   // The exception to this is if we have pre-received data. In that case there
1843   // is definitely data available immediately.
1844   if (preReceivedData_ && !preReceivedData_->empty()) {
1845     handleRead();
1846   }
1847 }
1848
1849 void AsyncSocket::handleInitialReadWrite() noexcept {
1850   // Our callers should already be holding a DestructorGuard, but grab
1851   // one here just to make sure, in case one of our calling code paths ever
1852   // changes.
1853   DestructorGuard dg(this);
1854   // If we have a readCallback_, make sure we enable read events.  We
1855   // may already be registered for reads if connectSuccess() set
1856   // the read calback.
1857   if (readCallback_ && !(eventFlags_ & EventHandler::READ)) {
1858     assert(state_ == StateEnum::ESTABLISHED);
1859     assert((shutdownFlags_ & SHUT_READ) == 0);
1860     if (!updateEventRegistration(EventHandler::READ, 0)) {
1861       assert(state_ == StateEnum::ERROR);
1862       return;
1863     }
1864     checkForImmediateRead();
1865   } else if (readCallback_ == nullptr) {
1866     // Unregister for read events.
1867     updateEventRegistration(0, EventHandler::READ);
1868   }
1869
1870   // If we have write requests pending, try to send them immediately.
1871   // Since we just finished accepting, there is a very good chance that we can
1872   // write without blocking.
1873   //
1874   // However, we only process them if EventHandler::WRITE is not already set,
1875   // which means that we're already blocked on a write attempt.  (This can
1876   // happen if connectSuccess() called write() before returning.)
1877   if (writeReqHead_ && !(eventFlags_ & EventHandler::WRITE)) {
1878     // Call handleWrite() to perform write processing.
1879     handleWrite();
1880   } else if (writeReqHead_ == nullptr) {
1881     // Unregister for write event.
1882     updateEventRegistration(0, EventHandler::WRITE);
1883   }
1884 }
1885
1886 void AsyncSocket::handleConnect() noexcept {
1887   VLOG(5) << "AsyncSocket::handleConnect() this=" << this << ", fd=" << fd_
1888           << ", state=" << state_;
1889   assert(state_ == StateEnum::CONNECTING);
1890   // SHUT_WRITE can never be set while we are still connecting;
1891   // SHUT_WRITE_PENDING may be set, be we only set SHUT_WRITE once the connect
1892   // finishes
1893   assert((shutdownFlags_ & SHUT_WRITE) == 0);
1894
1895   // In case we had a connect timeout, cancel the timeout
1896   writeTimeout_.cancelTimeout();
1897   // We don't use a persistent registration when waiting on a connect event,
1898   // so we have been automatically unregistered now.  Update eventFlags_ to
1899   // reflect reality.
1900   assert(eventFlags_ == EventHandler::WRITE);
1901   eventFlags_ = EventHandler::NONE;
1902
1903   // Call getsockopt() to check if the connect succeeded
1904   int error;
1905   socklen_t len = sizeof(error);
1906   int rv = getsockopt(fd_, SOL_SOCKET, SO_ERROR, &error, &len);
1907   if (rv != 0) {
1908     auto errnoCopy = errno;
1909     AsyncSocketException ex(
1910         AsyncSocketException::INTERNAL_ERROR,
1911         withAddr("error calling getsockopt() after connect"),
1912         errnoCopy);
1913     VLOG(4) << "AsyncSocket::handleConnect(this=" << this << ", fd="
1914                << fd_ << " host=" << addr_.describe()
1915                << ") exception:" << ex.what();
1916     return failConnect(__func__, ex);
1917   }
1918
1919   if (error != 0) {
1920     AsyncSocketException ex(AsyncSocketException::NOT_OPEN,
1921                            "connect failed", error);
1922     VLOG(1) << "AsyncSocket::handleConnect(this=" << this << ", fd="
1923             << fd_ << " host=" << addr_.describe()
1924             << ") exception: " << ex.what();
1925     return failConnect(__func__, ex);
1926   }
1927
1928   // Move into STATE_ESTABLISHED
1929   state_ = StateEnum::ESTABLISHED;
1930
1931   // If SHUT_WRITE_PENDING is set and we don't have any write requests to
1932   // perform, immediately shutdown the write half of the socket.
1933   if ((shutdownFlags_ & SHUT_WRITE_PENDING) && writeReqHead_ == nullptr) {
1934     // SHUT_READ shouldn't be set.  If close() is called on the socket while we
1935     // are still connecting we just abort the connect rather than waiting for
1936     // it to complete.
1937     assert((shutdownFlags_ & SHUT_READ) == 0);
1938     shutdown(fd_, SHUT_WR);
1939     shutdownFlags_ |= SHUT_WRITE;
1940   }
1941
1942   VLOG(7) << "AsyncSocket " << this << ": fd " << fd_
1943           << "successfully connected; state=" << state_;
1944
1945   // Remember the EventBase we are attached to, before we start invoking any
1946   // callbacks (since the callbacks may call detachEventBase()).
1947   EventBase* originalEventBase = eventBase_;
1948
1949   invokeConnectSuccess();
1950   // Note that the connect callback may have changed our state.
1951   // (set or unset the read callback, called write(), closed the socket, etc.)
1952   // The following code needs to handle these situations correctly.
1953   //
1954   // If the socket has been closed, readCallback_ and writeReqHead_ will
1955   // always be nullptr, so that will prevent us from trying to read or write.
1956   //
1957   // The main thing to check for is if eventBase_ is still originalEventBase.
1958   // If not, we have been detached from this event base, so we shouldn't
1959   // perform any more operations.
1960   if (eventBase_ != originalEventBase) {
1961     return;
1962   }
1963
1964   handleInitialReadWrite();
1965 }
1966
1967 void AsyncSocket::timeoutExpired() noexcept {
1968   VLOG(7) << "AsyncSocket " << this << ", fd " << fd_ << ": timeout expired: "
1969           << "state=" << state_ << ", events=" << std::hex << eventFlags_;
1970   DestructorGuard dg(this);
1971   assert(eventBase_->isInEventBaseThread());
1972
1973   if (state_ == StateEnum::CONNECTING) {
1974     // connect() timed out
1975     // Unregister for I/O events.
1976     if (connectCallback_) {
1977       AsyncSocketException ex(
1978           AsyncSocketException::TIMED_OUT,
1979           folly::sformat(
1980               "connect timed out after {}ms", connectTimeout_.count()));
1981       failConnect(__func__, ex);
1982     } else {
1983       // we faced a connect error without a connect callback, which could
1984       // happen due to TFO.
1985       AsyncSocketException ex(
1986           AsyncSocketException::TIMED_OUT, "write timed out during connection");
1987       failWrite(__func__, ex);
1988     }
1989   } else {
1990     // a normal write operation timed out
1991     AsyncSocketException ex(
1992         AsyncSocketException::TIMED_OUT,
1993         folly::sformat("write timed out after {}ms", sendTimeout_));
1994     failWrite(__func__, ex);
1995   }
1996 }
1997
1998 ssize_t AsyncSocket::tfoSendMsg(int fd, struct msghdr* msg, int msg_flags) {
1999   return detail::tfo_sendmsg(fd, msg, msg_flags);
2000 }
2001
2002 AsyncSocket::WriteResult
2003 AsyncSocket::sendSocketMessage(int fd, struct msghdr* msg, int msg_flags) {
2004   ssize_t totalWritten = 0;
2005   if (state_ == StateEnum::FAST_OPEN) {
2006     sockaddr_storage addr;
2007     auto len = addr_.getAddress(&addr);
2008     msg->msg_name = &addr;
2009     msg->msg_namelen = len;
2010     totalWritten = tfoSendMsg(fd_, msg, msg_flags);
2011     if (totalWritten >= 0) {
2012       tfoFinished_ = true;
2013       state_ = StateEnum::ESTABLISHED;
2014       // We schedule this asynchrously so that we don't end up
2015       // invoking initial read or write while a write is in progress.
2016       scheduleInitialReadWrite();
2017     } else if (errno == EINPROGRESS) {
2018       VLOG(4) << "TFO falling back to connecting";
2019       // A normal sendmsg doesn't return EINPROGRESS, however
2020       // TFO might fallback to connecting if there is no
2021       // cookie.
2022       state_ = StateEnum::CONNECTING;
2023       try {
2024         scheduleConnectTimeout();
2025         registerForConnectEvents();
2026       } catch (const AsyncSocketException& ex) {
2027         return WriteResult(
2028             WRITE_ERROR, std::make_unique<AsyncSocketException>(ex));
2029       }
2030       // Let's fake it that no bytes were written and return an errno.
2031       errno = EAGAIN;
2032       totalWritten = -1;
2033     } else if (errno == EOPNOTSUPP) {
2034       // Try falling back to connecting.
2035       VLOG(4) << "TFO not supported";
2036       state_ = StateEnum::CONNECTING;
2037       try {
2038         int ret = socketConnect((const sockaddr*)&addr, len);
2039         if (ret == 0) {
2040           // connect succeeded immediately
2041           // Treat this like no data was written.
2042           state_ = StateEnum::ESTABLISHED;
2043           scheduleInitialReadWrite();
2044         }
2045         // If there was no exception during connections,
2046         // we would return that no bytes were written.
2047         errno = EAGAIN;
2048         totalWritten = -1;
2049       } catch (const AsyncSocketException& ex) {
2050         return WriteResult(
2051             WRITE_ERROR, std::make_unique<AsyncSocketException>(ex));
2052       }
2053     } else if (errno == EAGAIN) {
2054       // Normally sendmsg would indicate that the write would block.
2055       // However in the fast open case, it would indicate that sendmsg
2056       // fell back to a connect. This is a return code from connect()
2057       // instead, and is an error condition indicating no fds available.
2058       return WriteResult(
2059           WRITE_ERROR,
2060           std::make_unique<AsyncSocketException>(
2061               AsyncSocketException::UNKNOWN, "No more free local ports"));
2062     }
2063   } else {
2064     totalWritten = ::sendmsg(fd, msg, msg_flags);
2065   }
2066   return WriteResult(totalWritten);
2067 }
2068
2069 AsyncSocket::WriteResult AsyncSocket::performWrite(
2070     const iovec* vec,
2071     uint32_t count,
2072     WriteFlags flags,
2073     uint32_t* countWritten,
2074     uint32_t* partialWritten) {
2075   // We use sendmsg() instead of writev() so that we can pass in MSG_NOSIGNAL
2076   // We correctly handle EPIPE errors, so we never want to receive SIGPIPE
2077   // (since it may terminate the program if the main program doesn't explicitly
2078   // ignore it).
2079   struct msghdr msg;
2080   msg.msg_name = nullptr;
2081   msg.msg_namelen = 0;
2082   msg.msg_iov = const_cast<iovec *>(vec);
2083   msg.msg_iovlen = std::min<size_t>(count, kIovMax);
2084   msg.msg_flags = 0;
2085   msg.msg_controllen = sendMsgParamCallback_->getAncillaryDataSize(flags);
2086   CHECK_GE(AsyncSocket::SendMsgParamsCallback::maxAncillaryDataSize,
2087            msg.msg_controllen);
2088
2089   if (msg.msg_controllen != 0) {
2090     msg.msg_control = reinterpret_cast<char*>(alloca(msg.msg_controllen));
2091     sendMsgParamCallback_->getAncillaryData(flags, msg.msg_control);
2092   } else {
2093     msg.msg_control = nullptr;
2094   }
2095   int msg_flags = sendMsgParamCallback_->getFlags(flags);
2096
2097   auto writeResult = sendSocketMessage(fd_, &msg, msg_flags);
2098   auto totalWritten = writeResult.writeReturn;
2099   if (totalWritten < 0) {
2100     bool tryAgain = (errno == EAGAIN);
2101 #ifdef __APPLE__
2102     // Apple has a bug where doing a second write on a socket which we
2103     // have opened with TFO causes an ENOTCONN to be thrown. However the
2104     // socket is really connected, so treat ENOTCONN as a EAGAIN until
2105     // this bug is fixed.
2106     tryAgain |= (errno == ENOTCONN);
2107 #endif
2108     if (!writeResult.exception && tryAgain) {
2109       // TCP buffer is full; we can't write any more data right now.
2110       *countWritten = 0;
2111       *partialWritten = 0;
2112       return WriteResult(0);
2113     }
2114     // error
2115     *countWritten = 0;
2116     *partialWritten = 0;
2117     return writeResult;
2118   }
2119
2120   appBytesWritten_ += totalWritten;
2121
2122   uint32_t bytesWritten;
2123   uint32_t n;
2124   for (bytesWritten = uint32_t(totalWritten), n = 0; n < count; ++n) {
2125     const iovec* v = vec + n;
2126     if (v->iov_len > bytesWritten) {
2127       // Partial write finished in the middle of this iovec
2128       *countWritten = n;
2129       *partialWritten = bytesWritten;
2130       return WriteResult(totalWritten);
2131     }
2132
2133     bytesWritten -= uint32_t(v->iov_len);
2134   }
2135
2136   assert(bytesWritten == 0);
2137   *countWritten = n;
2138   *partialWritten = 0;
2139   return WriteResult(totalWritten);
2140 }
2141
2142 /**
2143  * Re-register the EventHandler after eventFlags_ has changed.
2144  *
2145  * If an error occurs, fail() is called to move the socket into the error state
2146  * and call all currently installed callbacks.  After an error, the
2147  * AsyncSocket is completely unregistered.
2148  *
2149  * @return Returns true on succcess, or false on error.
2150  */
2151 bool AsyncSocket::updateEventRegistration() {
2152   VLOG(5) << "AsyncSocket::updateEventRegistration(this=" << this
2153           << ", fd=" << fd_ << ", evb=" << eventBase_ << ", state=" << state_
2154           << ", events=" << std::hex << eventFlags_;
2155   assert(eventBase_->isInEventBaseThread());
2156   if (eventFlags_ == EventHandler::NONE) {
2157     ioHandler_.unregisterHandler();
2158     return true;
2159   }
2160
2161   // Always register for persistent events, so we don't have to re-register
2162   // after being called back.
2163   if (!ioHandler_.registerHandler(
2164           uint16_t(eventFlags_ | EventHandler::PERSIST))) {
2165     eventFlags_ = EventHandler::NONE; // we're not registered after error
2166     AsyncSocketException ex(AsyncSocketException::INTERNAL_ERROR,
2167         withAddr("failed to update AsyncSocket event registration"));
2168     fail("updateEventRegistration", ex);
2169     return false;
2170   }
2171
2172   return true;
2173 }
2174
2175 bool AsyncSocket::updateEventRegistration(uint16_t enable,
2176                                            uint16_t disable) {
2177   uint16_t oldFlags = eventFlags_;
2178   eventFlags_ |= enable;
2179   eventFlags_ &= ~disable;
2180   if (eventFlags_ == oldFlags) {
2181     return true;
2182   } else {
2183     return updateEventRegistration();
2184   }
2185 }
2186
2187 void AsyncSocket::startFail() {
2188   // startFail() should only be called once
2189   assert(state_ != StateEnum::ERROR);
2190   assert(getDestructorGuardCount() > 0);
2191   state_ = StateEnum::ERROR;
2192   // Ensure that SHUT_READ and SHUT_WRITE are set,
2193   // so all future attempts to read or write will be rejected
2194   shutdownFlags_ |= (SHUT_READ | SHUT_WRITE);
2195
2196   if (eventFlags_ != EventHandler::NONE) {
2197     eventFlags_ = EventHandler::NONE;
2198     ioHandler_.unregisterHandler();
2199   }
2200   writeTimeout_.cancelTimeout();
2201
2202   if (fd_ >= 0) {
2203     ioHandler_.changeHandlerFD(-1);
2204     doClose();
2205   }
2206 }
2207
2208 void AsyncSocket::invokeAllErrors(const AsyncSocketException& ex) {
2209   invokeConnectErr(ex);
2210   failAllWrites(ex);
2211
2212   if (readCallback_) {
2213     ReadCallback* callback = readCallback_;
2214     readCallback_ = nullptr;
2215     callback->readErr(ex);
2216   }
2217 }
2218
2219 void AsyncSocket::finishFail() {
2220   assert(state_ == StateEnum::ERROR);
2221   assert(getDestructorGuardCount() > 0);
2222
2223   AsyncSocketException ex(
2224       AsyncSocketException::INTERNAL_ERROR,
2225       withAddr("socket closing after error"));
2226   invokeAllErrors(ex);
2227 }
2228
2229 void AsyncSocket::finishFail(const AsyncSocketException& ex) {
2230   assert(state_ == StateEnum::ERROR);
2231   assert(getDestructorGuardCount() > 0);
2232   invokeAllErrors(ex);
2233 }
2234
2235 void AsyncSocket::fail(const char* fn, const AsyncSocketException& ex) {
2236   VLOG(4) << "AsyncSocket(this=" << this << ", fd=" << fd_ << ", state="
2237              << state_ << " host=" << addr_.describe()
2238              << "): failed in " << fn << "(): "
2239              << ex.what();
2240   startFail();
2241   finishFail();
2242 }
2243
2244 void AsyncSocket::failConnect(const char* fn, const AsyncSocketException& ex) {
2245   VLOG(5) << "AsyncSocket(this=" << this << ", fd=" << fd_ << ", state="
2246                << state_ << " host=" << addr_.describe()
2247                << "): failed while connecting in " << fn << "(): "
2248                << ex.what();
2249   startFail();
2250
2251   invokeConnectErr(ex);
2252   finishFail(ex);
2253 }
2254
2255 void AsyncSocket::failRead(const char* fn, const AsyncSocketException& ex) {
2256   VLOG(5) << "AsyncSocket(this=" << this << ", fd=" << fd_ << ", state="
2257                << state_ << " host=" << addr_.describe()
2258                << "): failed while reading in " << fn << "(): "
2259                << ex.what();
2260   startFail();
2261
2262   if (readCallback_ != nullptr) {
2263     ReadCallback* callback = readCallback_;
2264     readCallback_ = nullptr;
2265     callback->readErr(ex);
2266   }
2267
2268   finishFail();
2269 }
2270
2271 void AsyncSocket::failErrMessageRead(const char* fn,
2272                                      const AsyncSocketException& ex) {
2273   VLOG(5) << "AsyncSocket(this=" << this << ", fd=" << fd_ << ", state="
2274                << state_ << " host=" << addr_.describe()
2275                << "): failed while reading message in " << fn << "(): "
2276                << ex.what();
2277   startFail();
2278
2279   if (errMessageCallback_ != nullptr) {
2280     ErrMessageCallback* callback = errMessageCallback_;
2281     errMessageCallback_ = nullptr;
2282     callback->errMessageError(ex);
2283   }
2284
2285   finishFail();
2286 }
2287
2288 void AsyncSocket::failWrite(const char* fn, const AsyncSocketException& ex) {
2289   VLOG(5) << "AsyncSocket(this=" << this << ", fd=" << fd_ << ", state="
2290                << state_ << " host=" << addr_.describe()
2291                << "): failed while writing in " << fn << "(): "
2292                << ex.what();
2293   startFail();
2294
2295   // Only invoke the first write callback, since the error occurred while
2296   // writing this request.  Let any other pending write callbacks be invoked in
2297   // finishFail().
2298   if (writeReqHead_ != nullptr) {
2299     WriteRequest* req = writeReqHead_;
2300     writeReqHead_ = req->getNext();
2301     WriteCallback* callback = req->getCallback();
2302     uint32_t bytesWritten = req->getTotalBytesWritten();
2303     req->destroy();
2304     if (callback) {
2305       callback->writeErr(bytesWritten, ex);
2306     }
2307   }
2308
2309   finishFail();
2310 }
2311
2312 void AsyncSocket::failWrite(const char* fn, WriteCallback* callback,
2313                              size_t bytesWritten,
2314                              const AsyncSocketException& ex) {
2315   // This version of failWrite() is used when the failure occurs before
2316   // we've added the callback to writeReqHead_.
2317   VLOG(4) << "AsyncSocket(this=" << this << ", fd=" << fd_ << ", state="
2318              << state_ << " host=" << addr_.describe()
2319              <<"): failed while writing in " << fn << "(): "
2320              << ex.what();
2321   startFail();
2322
2323   if (callback != nullptr) {
2324     callback->writeErr(bytesWritten, ex);
2325   }
2326
2327   finishFail();
2328 }
2329
2330 void AsyncSocket::failAllWrites(const AsyncSocketException& ex) {
2331   // Invoke writeError() on all write callbacks.
2332   // This is used when writes are forcibly shutdown with write requests
2333   // pending, or when an error occurs with writes pending.
2334   while (writeReqHead_ != nullptr) {
2335     WriteRequest* req = writeReqHead_;
2336     writeReqHead_ = req->getNext();
2337     WriteCallback* callback = req->getCallback();
2338     if (callback) {
2339       callback->writeErr(req->getTotalBytesWritten(), ex);
2340     }
2341     req->destroy();
2342   }
2343 }
2344
2345 void AsyncSocket::invalidState(ConnectCallback* callback) {
2346   VLOG(5) << "AsyncSocket(this=" << this << ", fd=" << fd_
2347           << "): connect() called in invalid state " << state_;
2348
2349   /*
2350    * The invalidState() methods don't use the normal failure mechanisms,
2351    * since we don't know what state we are in.  We don't want to call
2352    * startFail()/finishFail() recursively if we are already in the middle of
2353    * cleaning up.
2354    */
2355
2356   AsyncSocketException ex(AsyncSocketException::ALREADY_OPEN,
2357                          "connect() called with socket in invalid state");
2358   connectEndTime_ = std::chrono::steady_clock::now();
2359   if (state_ == StateEnum::CLOSED || state_ == StateEnum::ERROR) {
2360     if (callback) {
2361       callback->connectErr(ex);
2362     }
2363   } else {
2364     // We can't use failConnect() here since connectCallback_
2365     // may already be set to another callback.  Invoke this ConnectCallback
2366     // here; any other connectCallback_ will be invoked in finishFail()
2367     startFail();
2368     if (callback) {
2369       callback->connectErr(ex);
2370     }
2371     finishFail();
2372   }
2373 }
2374
2375 void AsyncSocket::invalidState(ErrMessageCallback* callback) {
2376   VLOG(4) << "AsyncSocket(this=" << this << ", fd=" << fd_
2377           << "): setErrMessageCB(" << callback
2378           << ") called in invalid state " << state_;
2379
2380   AsyncSocketException ex(
2381       AsyncSocketException::NOT_OPEN,
2382       msgErrQueueSupported
2383       ? "setErrMessageCB() called with socket in invalid state"
2384       : "This platform does not support socket error message notifications");
2385   if (state_ == StateEnum::CLOSED || state_ == StateEnum::ERROR) {
2386     if (callback) {
2387       callback->errMessageError(ex);
2388     }
2389   } else {
2390     startFail();
2391     if (callback) {
2392       callback->errMessageError(ex);
2393     }
2394     finishFail();
2395   }
2396 }
2397
2398 void AsyncSocket::invokeConnectErr(const AsyncSocketException& ex) {
2399   connectEndTime_ = std::chrono::steady_clock::now();
2400   if (connectCallback_) {
2401     ConnectCallback* callback = connectCallback_;
2402     connectCallback_ = nullptr;
2403     callback->connectErr(ex);
2404   }
2405 }
2406
2407 void AsyncSocket::invokeConnectSuccess() {
2408   connectEndTime_ = std::chrono::steady_clock::now();
2409   if (connectCallback_) {
2410     ConnectCallback* callback = connectCallback_;
2411     connectCallback_ = nullptr;
2412     callback->connectSuccess();
2413   }
2414 }
2415
2416 void AsyncSocket::invalidState(ReadCallback* callback) {
2417   VLOG(4) << "AsyncSocket(this=" << this << ", fd=" << fd_
2418              << "): setReadCallback(" << callback
2419              << ") called in invalid state " << state_;
2420
2421   AsyncSocketException ex(AsyncSocketException::NOT_OPEN,
2422                          "setReadCallback() called with socket in "
2423                          "invalid state");
2424   if (state_ == StateEnum::CLOSED || state_ == StateEnum::ERROR) {
2425     if (callback) {
2426       callback->readErr(ex);
2427     }
2428   } else {
2429     startFail();
2430     if (callback) {
2431       callback->readErr(ex);
2432     }
2433     finishFail();
2434   }
2435 }
2436
2437 void AsyncSocket::invalidState(WriteCallback* callback) {
2438   VLOG(4) << "AsyncSocket(this=" << this << ", fd=" << fd_
2439              << "): write() called in invalid state " << state_;
2440
2441   AsyncSocketException ex(AsyncSocketException::NOT_OPEN,
2442                          withAddr("write() called with socket in invalid state"));
2443   if (state_ == StateEnum::CLOSED || state_ == StateEnum::ERROR) {
2444     if (callback) {
2445       callback->writeErr(0, ex);
2446     }
2447   } else {
2448     startFail();
2449     if (callback) {
2450       callback->writeErr(0, ex);
2451     }
2452     finishFail();
2453   }
2454 }
2455
2456 void AsyncSocket::doClose() {
2457   if (fd_ == -1) return;
2458   if (shutdownSocketSet_) {
2459     shutdownSocketSet_->close(fd_);
2460   } else {
2461     ::close(fd_);
2462   }
2463   fd_ = -1;
2464 }
2465
2466 std::ostream& operator << (std::ostream& os,
2467                            const AsyncSocket::StateEnum& state) {
2468   os << static_cast<int>(state);
2469   return os;
2470 }
2471
2472 std::string AsyncSocket::withAddr(const std::string& s) {
2473   // Don't use addr_ directly because it may not be initialized
2474   // e.g. if constructed from fd
2475   folly::SocketAddress peer, local;
2476   try {
2477     getPeerAddress(&peer);
2478     getLocalAddress(&local);
2479   } catch (const std::exception&) {
2480     // ignore
2481   } catch (...) {
2482     // ignore
2483   }
2484   return s + " (peer=" + peer.describe() + ", local=" + local.describe() + ")";
2485 }
2486
2487 void AsyncSocket::setBufferCallback(BufferCallback* cb) {
2488   bufferCallback_ = cb;
2489 }
2490
2491 } // folly