From 5c74326fdc75ccdfc2152c15203625d8588096b6 Mon Sep 17 00:00:00 2001
From: Maxim Georgiev <maxgeorg@fb.com>
Date: Thu, 19 Jan 2017 16:14:24 -0800
Subject: [PATCH] Enable EOR flag configuration for folly::AsyncSocket.

Summary: EOR flag was hardcoded to "false" in folly::AsyncSocket. This diff enables changing EOR flag using AsyncSocket::setEorTracking() method.

Reviewed By: yfeldblum

Differential Revision: D4421966

fbshipit-source-id: 2e1b9b19ced6555845396ec33bfd3d5feb710640
---
 folly/io/async/AsyncSSLSocket.cpp        | 12 ++++--------
 folly/io/async/AsyncSSLSocket.h          |  4 ----
 folly/io/async/AsyncSocket.h             | 10 ++++++++--
 folly/io/async/test/AsyncSocketTest2.cpp |  7 +++++++
 4 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/folly/io/async/AsyncSSLSocket.cpp b/folly/io/async/AsyncSSLSocket.cpp
index 4b9f3173..c19a9326 100644
--- a/folly/io/async/AsyncSSLSocket.cpp
+++ b/folly/io/async/AsyncSSLSocket.cpp
@@ -358,13 +358,9 @@ std::string AsyncSSLSocket::getApplicationProtocol() noexcept {
   return "";
 }
 
-bool AsyncSSLSocket::isEorTrackingEnabled() const {
-  return trackEor_;
-}
-
 void AsyncSSLSocket::setEorTracking(bool track) {
-  if (trackEor_ != track) {
-    trackEor_ = track;
+  if (isEorTrackingEnabled() != track) {
+    AsyncSocket::setEorTracking(track);
     appEorByteNo_ = 0;
     minEorRawByteNo_ = 0;
   }
@@ -1538,7 +1534,7 @@ AsyncSocket::WriteResult AsyncSSLSocket::performWrite(
 
 int AsyncSSLSocket::eorAwareSSLWrite(SSL *ssl, const void *buf, int n,
                                       bool eor) {
-  if (eor && trackEor_) {
+  if (eor && isEorTrackingEnabled()) {
     if (appEorByteNo_) {
       // cannot track for more than one app byte EOR
       CHECK(appEorByteNo_ == appBytesWritten_ + n);
@@ -1601,7 +1597,7 @@ int AsyncSSLSocket::bioWrite(BIO* b, const char* in, int inl) {
   tsslSock = reinterpret_cast<AsyncSSLSocket*>(appData);
   CHECK(tsslSock);
 
-  if (tsslSock->trackEor_ && tsslSock->minEorRawByteNo_ &&
+  if (tsslSock->isEorTrackingEnabled() && tsslSock->minEorRawByteNo_ &&
       tsslSock->minEorRawByteNo_ <= BIO_number_written(b) + inl) {
     flags = MSG_EOR;
   }
diff --git a/folly/io/async/AsyncSSLSocket.h b/folly/io/async/AsyncSSLSocket.h
index e7d499e9..e9aca826 100644
--- a/folly/io/async/AsyncSSLSocket.h
+++ b/folly/io/async/AsyncSSLSocket.h
@@ -272,7 +272,6 @@ class AsyncSSLSocket : public virtual AsyncSocket {
 
   virtual std::string getSecurityProtocol() const override { return "TLS"; }
 
-  bool isEorTrackingEnabled() const override;
   virtual void setEorTracking(bool track) override;
   virtual size_t getRawBytesWritten() const override;
   virtual size_t getRawBytesReceived() const override;
@@ -792,9 +791,6 @@ class AsyncSSLSocket : public virtual AsyncSocket {
   // whether the SSL session was resumed using session ID or not
   bool sessionIDResumed_{false};
 
-  // Whether to track EOR or not.
-  bool trackEor_{false};
-
   // The app byte num that we are tracking for the MSG_EOR
   // Only one app EOR byte can be tracked.
   size_t appEorByteNo_{0};
diff --git a/folly/io/async/AsyncSocket.h b/folly/io/async/AsyncSocket.h
index e4f01a4d..5aeb159c 100644
--- a/folly/io/async/AsyncSocket.h
+++ b/folly/io/async/AsyncSocket.h
@@ -392,9 +392,13 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
   void getPeerAddress(
     folly::SocketAddress* address) const override;
 
-  bool isEorTrackingEnabled() const override { return false; }
+  bool isEorTrackingEnabled() const override {
+    return trackEor_;
+  }
 
-  void setEorTracking(bool /*track*/) override {}
+  void setEorTracking(bool track) override {
+    trackEor_ = track;
+  }
 
   bool connecting() const override {
     return (state_ == StateEnum::CONNECTING);
@@ -958,6 +962,8 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
   bool tfoAttempted_{false};
   bool tfoFinished_{false};
   bool noTransparentTls_{false};
+  // Whether to track EOR or not.
+  bool trackEor_{false};
 
   std::unique_ptr<EvbChangeCallback> evbChangeCb_{nullptr};
 };
diff --git a/folly/io/async/test/AsyncSocketTest2.cpp b/folly/io/async/test/AsyncSocketTest2.cpp
index b1f69db5..13f23b75 100644
--- a/folly/io/async/test/AsyncSocketTest2.cpp
+++ b/folly/io/async/test/AsyncSocketTest2.cpp
@@ -1177,6 +1177,13 @@ TEST(AsyncSocketTest, WriteIOBuf) {
   ReadCallback rcb;
   acceptedSocket->setReadCB(&rcb);
 
+  // Check if EOR tracking flag can be set and reset.
+  EXPECT_FALSE(socket->isEorTrackingEnabled());
+  socket->setEorTracking(true);
+  EXPECT_TRUE(socket->isEorTrackingEnabled());
+  socket->setEorTracking(false);
+  EXPECT_FALSE(socket->isEorTrackingEnabled());
+
   // Write a simple buffer to the socket
   constexpr size_t simpleBufLength = 5;
   char simpleBuf[simpleBufLength];
-- 
2.34.1