Make some AsyncTest methods virtual to allow mocking them using gtest/gmock
[folly.git] / folly / io / async / AsyncSocket.h
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 #pragma once
18
19 #include <folly/Optional.h>
20 #include <folly/SocketAddress.h>
21 #include <folly/detail/SocketFastOpen.h>
22 #include <folly/io/IOBuf.h>
23 #include <folly/io/ShutdownSocketSet.h>
24 #include <folly/io/async/AsyncSocketException.h>
25 #include <folly/io/async/AsyncTimeout.h>
26 #include <folly/io/async/AsyncTransport.h>
27 #include <folly/io/async/DelayedDestruction.h>
28 #include <folly/io/async/EventHandler.h>
29 #include <folly/portability/Sockets.h>
30
31 #include <sys/types.h>
32
33 #include <chrono>
34 #include <memory>
35 #include <map>
36
37 namespace folly {
38
39 /**
40  * A class for performing asynchronous I/O on a socket.
41  *
42  * AsyncSocket allows users to asynchronously wait for data on a socket, and
43  * to asynchronously send data.
44  *
45  * The APIs for reading and writing are intentionally asymmetric.  Waiting for
46  * data to read is a persistent API: a callback is installed, and is notified
47  * whenever new data is available.  It continues to be notified of new events
48  * until it is uninstalled.
49  *
50  * AsyncSocket does not provide read timeout functionality, because it
51  * typically cannot determine when the timeout should be active.  Generally, a
52  * timeout should only be enabled when processing is blocked waiting on data
53  * from the remote endpoint.  For server sockets, the timeout should not be
54  * active if the server is currently processing one or more outstanding
55  * requests for this socket.  For client sockets, the timeout should not be
56  * active if there are no requests pending on the socket.  Additionally, if a
57  * client has multiple pending requests, it will ususally want a separate
58  * timeout for each request, rather than a single read timeout.
59  *
60  * The write API is fairly intuitive: a user can request to send a block of
61  * data, and a callback will be informed once the entire block has been
62  * transferred to the kernel, or on error.  AsyncSocket does provide a send
63  * timeout, since most callers want to give up if the remote end stops
64  * responding and no further progress can be made sending the data.
65  */
66
67 #if defined __linux__ && !defined SO_NO_TRANSPARENT_TLS
68 #define SO_NO_TRANSPARENT_TLS 200
69 #endif
70
71 #ifdef _MSC_VER
72 // We do a dynamic_cast on this, in
73 // AsyncTransportWrapper::getUnderlyingTransport so be safe and
74 // force displacements for it. See:
75 // https://msdn.microsoft.com/en-us/library/7sf3txa8.aspx
76 #pragma vtordisp(push, 2)
77 #endif
78 class AsyncSocket : virtual public AsyncTransportWrapper {
79  public:
80   typedef std::unique_ptr<AsyncSocket, Destructor> UniquePtr;
81
82   class ConnectCallback {
83    public:
84     virtual ~ConnectCallback() = default;
85
86     /**
87      * connectSuccess() will be invoked when the connection has been
88      * successfully established.
89      */
90     virtual void connectSuccess() noexcept = 0;
91
92     /**
93      * connectErr() will be invoked if the connection attempt fails.
94      *
95      * @param ex        An exception describing the error that occurred.
96      */
97     virtual void connectErr(const AsyncSocketException& ex)
98       noexcept = 0;
99   };
100
101   class EvbChangeCallback {
102    public:
103     virtual ~EvbChangeCallback() = default;
104
105     // Called when the socket has been attached to a new EVB
106     // and is called from within that EVB thread
107     virtual void evbAttached(AsyncSocket* socket) = 0;
108
109     // Called when the socket is detached from an EVB and
110     // is called from the EVB thread being detached
111     virtual void evbDetached(AsyncSocket* socket) = 0;
112   };
113
114   /**
115    * This interface is implemented only for platforms supporting
116    * per-socket error queues.
117    */
118   class ErrMessageCallback {
119    public:
120     virtual ~ErrMessageCallback() = default;
121
122     /**
123      * errMessage() will be invoked when kernel puts a message to
124      * the error queue associated with the socket.
125      *
126      * @param cmsg      Reference to cmsghdr structure describing
127      *                  a message read from error queue associated
128      *                  with the socket.
129      */
130     virtual void
131     errMessage(const cmsghdr& cmsg) noexcept = 0;
132
133     /**
134      * errMessageError() will be invoked if an error occurs reading a message
135      * from the socket error stream.
136      *
137      * @param ex        An exception describing the error that occurred.
138      */
139     virtual void errMessageError(const AsyncSocketException& ex) noexcept = 0;
140   };
141
142   class SendMsgParamsCallback {
143    public:
144     virtual ~SendMsgParamsCallback() = default;
145
146     /**
147      * getFlags() will be invoked to retrieve the desired flags to be passed
148      * to ::sendmsg() system call. This method was intentionally declared
149      * non-virtual, so there is no way to override it. Instead feel free to
150      * override getFlagsImpl(flags, defaultFlags) method instead, and enjoy
151      * the convenience of defaultFlags passed there.
152      *
153      * @param flags     Write flags requested for the given write operation
154      */
155     int getFlags(folly::WriteFlags flags) noexcept {
156       return getFlagsImpl(flags, getDefaultFlags(flags));
157     }
158
159     /**
160      * getAncillaryData() will be invoked to initialize ancillary data
161      * buffer referred by "msg_control" field of msghdr structure passed to
162      * ::sendmsg() system call. The function assumes that the size of buffer
163      * is not smaller than the value returned by getAncillaryDataSize() method
164      * for the same combination of flags.
165      *
166      * @param flags     Write flags requested for the given write operation
167      * @param data      Pointer to ancillary data buffer to initialize.
168      */
169     virtual void getAncillaryData(
170       folly::WriteFlags /*flags*/,
171       void* /*data*/) noexcept {}
172
173     /**
174      * getAncillaryDataSize() will be invoked to retrieve the size of
175      * ancillary data buffer which should be passed to ::sendmsg() system call
176      *
177      * @param flags     Write flags requested for the given write operation
178      */
179     virtual uint32_t getAncillaryDataSize(folly::WriteFlags /*flags*/)
180         noexcept {
181       return 0;
182     }
183
184     static const size_t maxAncillaryDataSize{0x5000};
185
186    private:
187     /**
188      * getFlagsImpl() will be invoked by getFlags(folly::WriteFlags flags)
189      * method to retrieve the flags to be passed to ::sendmsg() system call.
190      * SendMsgParamsCallback::getFlags() is calling this method, and returns
191      * its results directly to the caller in AsyncSocket.
192      * Classes inheriting from SendMsgParamsCallback are welcome to override
193      * this method to force SendMsgParamsCallback to return its own set
194      * of flags.
195      *
196      * @param flags        Write flags requested for the given write operation
197      * @param defaultflags A set of message flags returned by getDefaultFlags()
198      *                     method for the given "flags" mask.
199      */
200     virtual int getFlagsImpl(folly::WriteFlags /*flags*/, int defaultFlags) {
201       return defaultFlags;
202     }
203
204     /**
205      * getDefaultFlags() will be invoked by  getFlags(folly::WriteFlags flags)
206      * to retrieve the default set of flags, and pass them to getFlagsImpl(...)
207      *
208      * @param flags     Write flags requested for the given write operation
209      */
210     int getDefaultFlags(folly::WriteFlags flags) noexcept;
211   };
212
213   explicit AsyncSocket();
214   /**
215    * Create a new unconnected AsyncSocket.
216    *
217    * connect() must later be called on this socket to establish a connection.
218    */
219   explicit AsyncSocket(EventBase* evb);
220
221   void setShutdownSocketSet(ShutdownSocketSet* ss);
222
223   /**
224    * Create a new AsyncSocket and begin the connection process.
225    *
226    * @param evb             EventBase that will manage this socket.
227    * @param address         The address to connect to.
228    * @param connectTimeout  Optional timeout in milliseconds for the connection
229    *                        attempt.
230    */
231   AsyncSocket(EventBase* evb,
232                const folly::SocketAddress& address,
233                uint32_t connectTimeout = 0);
234
235   /**
236    * Create a new AsyncSocket and begin the connection process.
237    *
238    * @param evb             EventBase that will manage this socket.
239    * @param ip              IP address to connect to (dotted-quad).
240    * @param port            Destination port in host byte order.
241    * @param connectTimeout  Optional timeout in milliseconds for the connection
242    *                        attempt.
243    */
244   AsyncSocket(EventBase* evb,
245                const std::string& ip,
246                uint16_t port,
247                uint32_t connectTimeout = 0);
248
249   /**
250    * Create a AsyncSocket from an already connected socket file descriptor.
251    *
252    * Note that while AsyncSocket enables TCP_NODELAY for sockets it creates
253    * when connecting, it does not change the socket options when given an
254    * existing file descriptor.  If callers want TCP_NODELAY enabled when using
255    * this version of the constructor, they need to explicitly call
256    * setNoDelay(true) after the constructor returns.
257    *
258    * @param evb EventBase that will manage this socket.
259    * @param fd  File descriptor to take over (should be a connected socket).
260    */
261   AsyncSocket(EventBase* evb, int fd);
262
263   /**
264    * Create an AsyncSocket from a different, already connected AsyncSocket.
265    *
266    * Similar to AsyncSocket(evb, fd) when fd was previously owned by an
267    * AsyncSocket.
268    */
269   explicit AsyncSocket(AsyncSocket::UniquePtr);
270
271   /**
272    * Helper function to create a shared_ptr<AsyncSocket>.
273    *
274    * This passes in the correct destructor object, since AsyncSocket's
275    * destructor is protected and cannot be invoked directly.
276    */
277   static std::shared_ptr<AsyncSocket> newSocket(EventBase* evb) {
278     return std::shared_ptr<AsyncSocket>(new AsyncSocket(evb),
279                                            Destructor());
280   }
281
282   /**
283    * Helper function to create a shared_ptr<AsyncSocket>.
284    */
285   static std::shared_ptr<AsyncSocket> newSocket(
286       EventBase* evb,
287       const folly::SocketAddress& address,
288       uint32_t connectTimeout = 0) {
289     return std::shared_ptr<AsyncSocket>(
290         new AsyncSocket(evb, address, connectTimeout),
291         Destructor());
292   }
293
294   /**
295    * Helper function to create a shared_ptr<AsyncSocket>.
296    */
297   static std::shared_ptr<AsyncSocket> newSocket(
298       EventBase* evb,
299       const std::string& ip,
300       uint16_t port,
301       uint32_t connectTimeout = 0) {
302     return std::shared_ptr<AsyncSocket>(
303         new AsyncSocket(evb, ip, port, connectTimeout),
304         Destructor());
305   }
306
307   /**
308    * Helper function to create a shared_ptr<AsyncSocket>.
309    */
310   static std::shared_ptr<AsyncSocket> newSocket(EventBase* evb, int fd) {
311     return std::shared_ptr<AsyncSocket>(new AsyncSocket(evb, fd),
312                                            Destructor());
313   }
314
315   /**
316    * Destroy the socket.
317    *
318    * AsyncSocket::destroy() must be called to destroy the socket.
319    * The normal destructor is private, and should not be invoked directly.
320    * This prevents callers from deleting a AsyncSocket while it is invoking a
321    * callback.
322    */
323   virtual void destroy() override;
324
325   /**
326    * Get the EventBase used by this socket.
327    */
328   EventBase* getEventBase() const override {
329     return eventBase_;
330   }
331
332   /**
333    * Get the file descriptor used by the AsyncSocket.
334    */
335   virtual int getFd() const {
336     return fd_;
337   }
338
339   /**
340    * Extract the file descriptor from the AsyncSocket.
341    *
342    * This will immediately cause any installed callbacks to be invoked with an
343    * error.  The AsyncSocket may no longer be used after the file descriptor
344    * has been extracted.
345    *
346    * This method should be used with care as the resulting fd is not guaranteed
347    * to perfectly reflect the state of the AsyncSocket (security state,
348    * pre-received data, etc.).
349    *
350    * Returns the file descriptor.  The caller assumes ownership of the
351    * descriptor, and it will not be closed when the AsyncSocket is destroyed.
352    */
353   virtual int detachFd();
354
355   /**
356    * Uniquely identifies a handle to a socket option value. Each
357    * combination of level and option name corresponds to one socket
358    * option value.
359    */
360   class OptionKey {
361    public:
362     bool operator<(const OptionKey& other) const {
363       if (level == other.level) {
364         return optname < other.optname;
365       }
366       return level < other.level;
367     }
368     int apply(int fd, int val) const {
369       return setsockopt(fd, level, optname, &val, sizeof(val));
370     }
371     int level;
372     int optname;
373   };
374
375   // Maps from a socket option key to its value
376   typedef std::map<OptionKey, int> OptionMap;
377
378   static const OptionMap emptyOptionMap;
379   static const folly::SocketAddress& anyAddress();
380
381   /**
382    * Initiate a connection.
383    *
384    * @param callback  The callback to inform when the connection attempt
385    *                  completes.
386    * @param address   The address to connect to.
387    * @param timeout   A timeout value, in milliseconds.  If the connection
388    *                  does not succeed within this period,
389    *                  callback->connectError() will be invoked.
390    */
391   virtual void connect(
392       ConnectCallback* callback,
393       const folly::SocketAddress& address,
394       int timeout = 0,
395       const OptionMap& options = emptyOptionMap,
396       const folly::SocketAddress& bindAddr = anyAddress()) noexcept;
397
398   void connect(
399       ConnectCallback* callback,
400       const std::string& ip,
401       uint16_t port,
402       int timeout = 0,
403       const OptionMap& options = emptyOptionMap) noexcept;
404
405   /**
406    * If a connect request is in-flight, cancels it and closes the socket
407    * immediately. Otherwise, this is a no-op.
408    *
409    * This does not invoke any connection related callbacks. Call this to
410    * prevent any connect callback while cleaning up, etc.
411    */
412   void cancelConnect();
413
414   /**
415    * Set the send timeout.
416    *
417    * If write requests do not make any progress for more than the specified
418    * number of milliseconds, fail all pending writes and close the socket.
419    *
420    * If write requests are currently pending when setSendTimeout() is called,
421    * the timeout interval is immediately restarted using the new value.
422    *
423    * (See the comments for AsyncSocket for an explanation of why AsyncSocket
424    * provides setSendTimeout() but not setRecvTimeout().)
425    *
426    * @param milliseconds  The timeout duration, in milliseconds.  If 0, no
427    *                      timeout will be used.
428    */
429   void setSendTimeout(uint32_t milliseconds) override;
430
431   /**
432    * Get the send timeout.
433    *
434    * @return Returns the current send timeout, in milliseconds.  A return value
435    *         of 0 indicates that no timeout is set.
436    */
437   uint32_t getSendTimeout() const override {
438     return sendTimeout_;
439   }
440
441   /**
442    * Set the maximum number of reads to execute from the underlying
443    * socket each time the EventBase detects that new ingress data is
444    * available. The default is unlimited, but callers can use this method
445    * to limit the amount of data read from the socket per event loop
446    * iteration.
447    *
448    * @param maxReads  Maximum number of reads per data-available event;
449    *                  a value of zero means unlimited.
450    */
451   void setMaxReadsPerEvent(uint16_t maxReads) {
452     maxReadsPerEvent_ = maxReads;
453   }
454
455   /**
456    * Get the maximum number of reads this object will execute from
457    * the underlying socket each time the EventBase detects that new
458    * ingress data is available.
459    *
460    * @returns Maximum number of reads per data-available event; a value
461    *          of zero means unlimited.
462    */
463   uint16_t getMaxReadsPerEvent() const {
464     return maxReadsPerEvent_;
465   }
466
467   /**
468    * Set a pointer to ErrMessageCallback implementation which will be
469    * receiving notifications for messages posted to the error queue
470    * associated with the socket.
471    * ErrMessageCallback is implemented only for platforms with
472    * per-socket error message queus support (recvmsg() system call must
473    * )
474    *
475    */
476   virtual void setErrMessageCB(ErrMessageCallback* callback);
477
478   /**
479    * Get a pointer to ErrMessageCallback implementation currently
480    * registered with this socket.
481    *
482    */
483   virtual ErrMessageCallback* getErrMessageCallback() const;
484
485   /**
486    * Set a pointer to SendMsgParamsCallback implementation which
487    * will be used to form ::sendmsg() system call parameters
488    *
489    */
490   virtual void setSendMsgParamCB(SendMsgParamsCallback* callback);
491
492   /**
493    * Get a pointer to SendMsgParamsCallback implementation currently
494    * registered with this socket.
495    *
496    */
497   virtual SendMsgParamsCallback* getSendMsgParamsCB() const;
498
499   // Read and write methods
500   void setReadCB(ReadCallback* callback) override;
501   ReadCallback* getReadCallback() const override;
502
503   void write(WriteCallback* callback, const void* buf, size_t bytes,
504              WriteFlags flags = WriteFlags::NONE) override;
505   void writev(WriteCallback* callback, const iovec* vec, size_t count,
506               WriteFlags flags = WriteFlags::NONE) override;
507   void writeChain(WriteCallback* callback,
508                   std::unique_ptr<folly::IOBuf>&& buf,
509                   WriteFlags flags = WriteFlags::NONE) override;
510
511   class WriteRequest;
512   virtual void writeRequest(WriteRequest* req);
513   void writeRequestReady() {
514     handleWrite();
515   }
516
517   // Methods inherited from AsyncTransport
518   void close() override;
519   void closeNow() override;
520   void closeWithReset() override;
521   void shutdownWrite() override;
522   void shutdownWriteNow() override;
523
524   bool readable() const override;
525   bool isPending() const override;
526   virtual bool hangup() const;
527   bool good() const override;
528   bool error() const override;
529   void attachEventBase(EventBase* eventBase) override;
530   void detachEventBase() override;
531   bool isDetachable() const override;
532
533   void getLocalAddress(
534     folly::SocketAddress* address) const override;
535   void getPeerAddress(
536     folly::SocketAddress* address) const override;
537
538   bool isEorTrackingEnabled() const override {
539     return trackEor_;
540   }
541
542   void setEorTracking(bool track) override {
543     trackEor_ = track;
544   }
545
546   bool connecting() const override {
547     return (state_ == StateEnum::CONNECTING);
548   }
549
550   virtual bool isClosedByPeer() const {
551     return (state_ == StateEnum::CLOSED &&
552             (readErr_ == READ_EOF || readErr_ == READ_ERROR));
553   }
554
555   virtual bool isClosedBySelf() const {
556     return (state_ == StateEnum::CLOSED &&
557             (readErr_ != READ_EOF && readErr_ != READ_ERROR));
558   }
559
560   size_t getAppBytesWritten() const override {
561     return appBytesWritten_;
562   }
563
564   size_t getRawBytesWritten() const override {
565     return getAppBytesWritten();
566   }
567
568   size_t getAppBytesReceived() const override {
569     return appBytesReceived_;
570   }
571
572   size_t getRawBytesReceived() const override {
573     return getAppBytesReceived();
574   }
575
576   std::chrono::nanoseconds getConnectTime() const {
577     return connectEndTime_ - connectStartTime_;
578   }
579
580   std::chrono::milliseconds getConnectTimeout() const {
581     return connectTimeout_;
582   }
583
584   bool getTFOAttempted() const {
585     return tfoAttempted_;
586   }
587
588   /**
589    * Returns whether or not the attempt to use TFO
590    * finished successfully. This does not necessarily
591    * mean TFO worked, just that trying to use TFO
592    * succeeded.
593    */
594   bool getTFOFinished() const {
595     return tfoFinished_;
596   }
597
598   /**
599    * Returns whether or not TFO attempt succeded on this
600    * connection.
601    * For servers this is pretty straightforward API and can
602    * be invoked right after the connection is accepted. This API
603    * will perform one syscall.
604    * This API is a bit tricky to use for clients, since clients
605    * only know this for sure after the SYN-ACK is returned. So it's
606    * appropriate to call this only after the first application
607    * data is read from the socket when the caller knows that
608    * the SYN has been ACKed by the server.
609    */
610   bool getTFOSucceded() const;
611
612   // Methods controlling socket options
613
614   /**
615    * Force writes to be transmitted immediately.
616    *
617    * This controls the TCP_NODELAY socket option.  When enabled, TCP segments
618    * are sent as soon as possible, even if it is not a full frame of data.
619    * When disabled, the data may be buffered briefly to try and wait for a full
620    * frame of data.
621    *
622    * By default, TCP_NODELAY is enabled for AsyncSocket objects.
623    *
624    * This method will fail if the socket is not currently open.
625    *
626    * @return Returns 0 if the TCP_NODELAY flag was successfully updated,
627    *         or a non-zero errno value on error.
628    */
629   int setNoDelay(bool noDelay);
630
631
632   /**
633    * Set the FD_CLOEXEC flag so that the socket will be closed if the program
634    * later forks and execs.
635    */
636   void setCloseOnExec();
637
638   /*
639    * Set the Flavor of Congestion Control to be used for this Socket
640    * Please check '/lib/modules/<kernel>/kernel/net/ipv4' for tcp_*.ko
641    * first to make sure the module is available for plugging in
642    * Alternatively you can choose from net.ipv4.tcp_allowed_congestion_control
643    */
644   int setCongestionFlavor(const std::string &cname);
645
646   /*
647    * Forces ACKs to be sent immediately
648    *
649    * @return Returns 0 if the TCP_QUICKACK flag was successfully updated,
650    *         or a non-zero errno value on error.
651    */
652   int setQuickAck(bool quickack);
653
654   /**
655    * Set the send bufsize
656    */
657   int setSendBufSize(size_t bufsize);
658
659   /**
660    * Set the recv bufsize
661    */
662   int setRecvBufSize(size_t bufsize);
663
664   /**
665    * Sets a specific tcp personality
666    * Available only on kernels 3.2 and greater
667    */
668   #define SO_SET_NAMESPACE        41
669   int setTCPProfile(int profd);
670
671   /**
672    * Generic API for reading a socket option.
673    *
674    * @param level     same as the "level" parameter in getsockopt().
675    * @param optname   same as the "optname" parameter in getsockopt().
676    * @param optval    pointer to the variable in which the option value should
677    *                  be returned.
678    * @param optlen    value-result argument, initially containing the size of
679    *                  the buffer pointed to by optval, and modified on return
680    *                  to indicate the actual size of the value returned.
681    * @return          same as the return value of getsockopt().
682    */
683   template <typename T>
684   int getSockOpt(int level, int optname, T* optval, socklen_t* optlen) {
685     return getsockopt(fd_, level, optname, (void*) optval, optlen);
686   }
687
688   /**
689    * Generic API for setting a socket option.
690    *
691    * @param level     same as the "level" parameter in getsockopt().
692    * @param optname   same as the "optname" parameter in getsockopt().
693    * @param optval    the option value to set.
694    * @return          same as the return value of setsockopt().
695    */
696   template <typename T>
697   int setSockOpt(int  level,  int  optname,  const T *optval) {
698     return setsockopt(fd_, level, optname, optval, sizeof(T));
699   }
700
701   /**
702    * Virtual method for reading a socket option returning integer
703    * value, which is the most typical case. Convenient for overriding
704    * and mocking.
705    *
706    * @param level     same as the "level" parameter in getsockopt().
707    * @param optname   same as the "optname" parameter in getsockopt().
708    * @param optval    same as "optval" parameter in getsockopt().
709    * @param optlen    same as "optlen" parameter in getsockopt().
710    * @return          same as the return value of getsockopt().
711    */
712   virtual int
713   getSockOptVirtual(int level, int optname, void* optval, socklen_t* optlen) {
714     return getsockopt(fd_, level, optname, optval, optlen);
715   }
716
717   /**
718    * Virtual method for setting a socket option accepting integer
719    * value, which is the most typical case. Convenient for overriding
720    * and mocking.
721    *
722    * @param level     same as the "level" parameter in setsockopt().
723    * @param optname   same as the "optname" parameter in setsockopt().
724    * @param optval    same as "optval" parameter in setsockopt().
725    * @param optlen    same as "optlen" parameter in setsockopt().
726    * @return          same as the return value of setsockopt().
727    */
728   virtual int setSockOptVirtual(
729       int level,
730       int optname,
731       void const* optval,
732       socklen_t optlen) {
733     return setsockopt(fd_, level, optname, optval, optlen);
734   }
735
736   /**
737    * Set pre-received data, to be returned to read callback before any data
738    * from the socket.
739    */
740   virtual void setPreReceivedData(std::unique_ptr<IOBuf> data) {
741     if (preReceivedData_) {
742       preReceivedData_->prependChain(std::move(data));
743     } else {
744       preReceivedData_ = std::move(data);
745     }
746   }
747
748   /**
749    * Enables TFO behavior on the AsyncSocket if FOLLY_ALLOW_TFO
750    * is set.
751    */
752   void enableTFO() {
753     // No-op if folly does not allow tfo
754 #if FOLLY_ALLOW_TFO
755     tfoEnabled_ = true;
756 #endif
757   }
758
759   void disableTransparentTls() {
760     noTransparentTls_ = true;
761   }
762
763   enum class StateEnum : uint8_t {
764     UNINIT,
765     CONNECTING,
766     ESTABLISHED,
767     CLOSED,
768     ERROR,
769     FAST_OPEN,
770   };
771
772   void setBufferCallback(BufferCallback* cb);
773
774   // Callers should set this prior to connecting the socket for the safest
775   // behavior.
776   void setEvbChangedCallback(std::unique_ptr<EvbChangeCallback> cb) {
777     evbChangeCb_ = std::move(cb);
778   }
779
780   /**
781    * writeReturn is the total number of bytes written, or WRITE_ERROR on error.
782    * If no data has been written, 0 is returned.
783    * exception is a more specific exception that cause a write error.
784    * Not all writes have exceptions associated with them thus writeReturn
785    * should be checked to determine whether the operation resulted in an error.
786    */
787   struct WriteResult {
788     explicit WriteResult(ssize_t ret) : writeReturn(ret) {}
789
790     WriteResult(ssize_t ret, std::unique_ptr<const AsyncSocketException> e)
791         : writeReturn(ret), exception(std::move(e)) {}
792
793     ssize_t writeReturn;
794     std::unique_ptr<const AsyncSocketException> exception;
795   };
796
797   /**
798    * readReturn is the number of bytes read, or READ_EOF on EOF, or
799    * READ_ERROR on error, or READ_BLOCKING if the operation will
800    * block.
801    * exception is a more specific exception that may have caused a read error.
802    * Not all read errors have exceptions associated with them thus readReturn
803    * should be checked to determine whether the operation resulted in an error.
804    */
805   struct ReadResult {
806     explicit ReadResult(ssize_t ret) : readReturn(ret) {}
807
808     ReadResult(ssize_t ret, std::unique_ptr<const AsyncSocketException> e)
809         : readReturn(ret), exception(std::move(e)) {}
810
811     ssize_t readReturn;
812     std::unique_ptr<const AsyncSocketException> exception;
813   };
814
815   /**
816    * A WriteRequest object tracks information about a pending write operation.
817    */
818   class WriteRequest {
819    public:
820     WriteRequest(AsyncSocket* socket, WriteCallback* callback) :
821       socket_(socket), callback_(callback) {}
822
823     virtual void start() {}
824
825     virtual void destroy() = 0;
826
827     virtual WriteResult performWrite() = 0;
828
829     virtual void consume() = 0;
830
831     virtual bool isComplete() = 0;
832
833     WriteRequest* getNext() const {
834       return next_;
835     }
836
837     WriteCallback* getCallback() const {
838       return callback_;
839     }
840
841     uint32_t getTotalBytesWritten() const {
842       return totalBytesWritten_;
843     }
844
845     void append(WriteRequest* next) {
846       assert(next_ == nullptr);
847       next_ = next;
848     }
849
850     void fail(const char* fn, const AsyncSocketException& ex) {
851       socket_->failWrite(fn, ex);
852     }
853
854     void bytesWritten(size_t count) {
855       totalBytesWritten_ += uint32_t(count);
856       socket_->appBytesWritten_ += count;
857     }
858
859    protected:
860     // protected destructor, to ensure callers use destroy()
861     virtual ~WriteRequest() {}
862
863     AsyncSocket* socket_;         ///< parent socket
864     WriteRequest* next_{nullptr};          ///< pointer to next WriteRequest
865     WriteCallback* callback_;     ///< completion callback
866     uint32_t totalBytesWritten_{0};  ///< total bytes written
867   };
868
869  protected:
870   enum ReadResultEnum {
871     READ_EOF = 0,
872     READ_ERROR = -1,
873     READ_BLOCKING = -2,
874     READ_NO_ERROR = -3,
875   };
876
877   enum WriteResultEnum {
878     WRITE_ERROR = -1,
879   };
880
881   /**
882    * Protected destructor.
883    *
884    * Users of AsyncSocket must never delete it directly.  Instead, invoke
885    * destroy() instead.  (See the documentation in DelayedDestruction.h for
886    * more details.)
887    */
888   ~AsyncSocket();
889
890   friend std::ostream& operator << (std::ostream& os, const StateEnum& state);
891
892   enum ShutdownFlags {
893     /// shutdownWrite() called, but we are still waiting on writes to drain
894     SHUT_WRITE_PENDING = 0x01,
895     /// writes have been completely shut down
896     SHUT_WRITE = 0x02,
897     /**
898      * Reads have been shutdown.
899      *
900      * At the moment we don't distinguish between remote read shutdown
901      * (received EOF from the remote end) and local read shutdown.  We can
902      * only receive EOF when a read callback is set, and we immediately inform
903      * it of the EOF.  Therefore there doesn't seem to be any reason to have a
904      * separate state of "received EOF but the local side may still want to
905      * read".
906      *
907      * We also don't currently provide any API for only shutting down the read
908      * side of a socket.  (This is a no-op as far as TCP is concerned, anyway.)
909      */
910     SHUT_READ = 0x04,
911   };
912
913   class BytesWriteRequest;
914
915   class WriteTimeout : public AsyncTimeout {
916    public:
917     WriteTimeout(AsyncSocket* socket, EventBase* eventBase)
918       : AsyncTimeout(eventBase)
919       , socket_(socket) {}
920
921     virtual void timeoutExpired() noexcept {
922       socket_->timeoutExpired();
923     }
924
925    private:
926     AsyncSocket* socket_;
927   };
928
929   class IoHandler : public EventHandler {
930    public:
931     IoHandler(AsyncSocket* socket, EventBase* eventBase)
932       : EventHandler(eventBase, -1)
933       , socket_(socket) {}
934     IoHandler(AsyncSocket* socket, EventBase* eventBase, int fd)
935       : EventHandler(eventBase, fd)
936       , socket_(socket) {}
937
938     virtual void handlerReady(uint16_t events) noexcept {
939       socket_->ioReady(events);
940     }
941
942    private:
943     AsyncSocket* socket_;
944   };
945
946   void init();
947
948   class ImmediateReadCB : public folly::EventBase::LoopCallback {
949    public:
950     explicit ImmediateReadCB(AsyncSocket* socket) : socket_(socket) {}
951     void runLoopCallback() noexcept override {
952       DestructorGuard dg(socket_);
953       socket_->checkForImmediateRead();
954     }
955    private:
956     AsyncSocket* socket_;
957   };
958
959   /**
960    * Schedule checkForImmediateRead to be executed in the next loop
961    * iteration.
962    */
963   void scheduleImmediateRead() noexcept {
964     if (good()) {
965       eventBase_->runInLoop(&immediateReadHandler_);
966     }
967   }
968
969   /**
970    * Schedule handleInitalReadWrite to run in the next iteration.
971    */
972   void scheduleInitialReadWrite() noexcept {
973     if (good()) {
974       DestructorGuard dg(this);
975       eventBase_->runInLoop([this, dg] {
976         if (good()) {
977           handleInitialReadWrite();
978         }
979       });
980     }
981   }
982
983   // event notification methods
984   void ioReady(uint16_t events) noexcept;
985   virtual void checkForImmediateRead() noexcept;
986   virtual void handleInitialReadWrite() noexcept;
987   virtual void prepareReadBuffer(void** buf, size_t* buflen);
988   virtual void handleErrMessages() noexcept;
989   virtual void handleRead() noexcept;
990   virtual void handleWrite() noexcept;
991   virtual void handleConnect() noexcept;
992   void timeoutExpired() noexcept;
993
994   /**
995    * Attempt to read from the socket.
996    *
997    * @param buf      The buffer to read data into.
998    * @param buflen   The length of the buffer.
999    *
1000    * @return Returns a read result. See read result for details.
1001    */
1002   virtual ReadResult performRead(void** buf, size_t* buflen, size_t* offset);
1003
1004   /**
1005    * Populate an iovec array from an IOBuf and attempt to write it.
1006    *
1007    * @param callback Write completion/error callback.
1008    * @param vec      Target iovec array; caller retains ownership.
1009    * @param count    Number of IOBufs to write, beginning at start of buf.
1010    * @param buf      Chain of iovecs.
1011    * @param flags    set of flags for the underlying write calls, like cork
1012    */
1013   void writeChainImpl(WriteCallback* callback, iovec* vec,
1014                       size_t count, std::unique_ptr<folly::IOBuf>&& buf,
1015                       WriteFlags flags);
1016
1017   /**
1018    * Write as much data as possible to the socket without blocking,
1019    * and queue up any leftover data to send when the socket can
1020    * handle writes again.
1021    *
1022    * @param callback The callback to invoke when the write is completed.
1023    * @param vec      Array of buffers to write; this method will make a
1024    *                 copy of the vector (but not the buffers themselves)
1025    *                 if the write has to be completed asynchronously.
1026    * @param count    Number of elements in vec.
1027    * @param buf      The IOBuf that manages the buffers referenced by
1028    *                 vec, or a pointer to nullptr if the buffers are not
1029    *                 associated with an IOBuf.  Note that ownership of
1030    *                 the IOBuf is transferred here; upon completion of
1031    *                 the write, the AsyncSocket deletes the IOBuf.
1032    * @param flags    Set of write flags.
1033    */
1034   void writeImpl(WriteCallback* callback, const iovec* vec, size_t count,
1035                  std::unique_ptr<folly::IOBuf>&& buf,
1036                  WriteFlags flags = WriteFlags::NONE);
1037
1038   /**
1039    * Attempt to write to the socket.
1040    *
1041    * @param vec             The iovec array pointing to the buffers to write.
1042    * @param count           The length of the iovec array.
1043    * @param flags           Set of write flags.
1044    * @param countWritten    On return, the value pointed to by this parameter
1045    *                          will contain the number of iovec entries that were
1046    *                          fully written.
1047    * @param partialWritten  On return, the value pointed to by this parameter
1048    *                          will contain the number of bytes written in the
1049    *                          partially written iovec entry.
1050    *
1051    * @return Returns a WriteResult. See WriteResult for more details.
1052    */
1053   virtual WriteResult performWrite(
1054       const iovec* vec,
1055       uint32_t count,
1056       WriteFlags flags,
1057       uint32_t* countWritten,
1058       uint32_t* partialWritten);
1059
1060   /**
1061    * Sends the message over the socket using sendmsg
1062    *
1063    * @param msg       Message to send
1064    * @param msg_flags Flags to pass to sendmsg
1065    */
1066   AsyncSocket::WriteResult
1067   sendSocketMessage(int fd, struct msghdr* msg, int msg_flags);
1068
1069   virtual ssize_t tfoSendMsg(int fd, struct msghdr* msg, int msg_flags);
1070
1071   int socketConnect(const struct sockaddr* addr, socklen_t len);
1072
1073   virtual void scheduleConnectTimeout();
1074   void registerForConnectEvents();
1075
1076   bool updateEventRegistration();
1077
1078   /**
1079    * Update event registration.
1080    *
1081    * @param enable Flags of events to enable. Set it to 0 if no events
1082    * need to be enabled in this call.
1083    * @param disable Flags of events
1084    * to disable. Set it to 0 if no events need to be disabled in this
1085    * call.
1086    *
1087    * @return true iff the update is successful.
1088    */
1089   bool updateEventRegistration(uint16_t enable, uint16_t disable);
1090
1091   // Actually close the file descriptor and set it to -1 so we don't
1092   // accidentally close it again.
1093   void doClose();
1094
1095   // error handling methods
1096   void startFail();
1097   void finishFail();
1098   void finishFail(const AsyncSocketException& ex);
1099   void invokeAllErrors(const AsyncSocketException& ex);
1100   void fail(const char* fn, const AsyncSocketException& ex);
1101   void failConnect(const char* fn, const AsyncSocketException& ex);
1102   void failRead(const char* fn, const AsyncSocketException& ex);
1103   void failErrMessageRead(const char* fn, const AsyncSocketException& ex);
1104   void failWrite(const char* fn, WriteCallback* callback, size_t bytesWritten,
1105                  const AsyncSocketException& ex);
1106   void failWrite(const char* fn, const AsyncSocketException& ex);
1107   void failAllWrites(const AsyncSocketException& ex);
1108   virtual void invokeConnectErr(const AsyncSocketException& ex);
1109   virtual void invokeConnectSuccess();
1110   void invalidState(ConnectCallback* callback);
1111   void invalidState(ErrMessageCallback* callback);
1112   void invalidState(ReadCallback* callback);
1113   void invalidState(WriteCallback* callback);
1114
1115   std::string withAddr(const std::string& s);
1116
1117   StateEnum state_;                      ///< StateEnum describing current state
1118   uint8_t shutdownFlags_;                ///< Shutdown state (ShutdownFlags)
1119   uint16_t eventFlags_;                  ///< EventBase::HandlerFlags settings
1120   int fd_;                               ///< The socket file descriptor
1121   mutable folly::SocketAddress addr_;    ///< The address we tried to connect to
1122   mutable folly::SocketAddress localAddr_;
1123                                          ///< The address we are connecting from
1124   uint32_t sendTimeout_;                 ///< The send timeout, in milliseconds
1125   uint16_t maxReadsPerEvent_;            ///< Max reads per event loop iteration
1126   EventBase* eventBase_;                 ///< The EventBase
1127   WriteTimeout writeTimeout_;            ///< A timeout for connect and write
1128   IoHandler ioHandler_;                  ///< A EventHandler to monitor the fd
1129   ImmediateReadCB immediateReadHandler_; ///< LoopCallback for checking read
1130
1131   ConnectCallback* connectCallback_;     ///< ConnectCallback
1132   ErrMessageCallback* errMessageCallback_; ///< TimestampCallback
1133   SendMsgParamsCallback*                 ///< Callback for retreaving
1134       sendMsgParamCallback_;             ///< ::sendmsg() parameters
1135   ReadCallback* readCallback_;           ///< ReadCallback
1136   WriteRequest* writeReqHead_;           ///< Chain of WriteRequests
1137   WriteRequest* writeReqTail_;           ///< End of WriteRequest chain
1138   ShutdownSocketSet* shutdownSocketSet_;
1139   size_t appBytesReceived_;              ///< Num of bytes received from socket
1140   size_t appBytesWritten_;               ///< Num of bytes written to socket
1141   bool isBufferMovable_{false};
1142
1143   // Pre-received data, to be returned to read callback before any data from the
1144   // socket.
1145   std::unique_ptr<IOBuf> preReceivedData_;
1146
1147   int8_t readErr_{READ_NO_ERROR};        ///< The read error encountered, if any
1148
1149   std::chrono::steady_clock::time_point connectStartTime_;
1150   std::chrono::steady_clock::time_point connectEndTime_;
1151
1152   std::chrono::milliseconds connectTimeout_{0};
1153
1154   BufferCallback* bufferCallback_{nullptr};
1155   bool tfoEnabled_{false};
1156   bool tfoAttempted_{false};
1157   bool tfoFinished_{false};
1158   bool noTransparentTls_{false};
1159   // Whether to track EOR or not.
1160   bool trackEor_{false};
1161
1162   std::unique_ptr<EvbChangeCallback> evbChangeCb_{nullptr};
1163 };
1164 #ifdef _MSC_VER
1165 #pragma vtordisp(pop)
1166 #endif
1167
1168 } // folly