Os: add ReceiveDatagram am: d75fd9e9b7 am: da734f5096

am: 298f6b81b5

Change-Id: Ifb2ea443615b2fb75c024cb0c1f6b4919793ba69

7 files changed, 135 insertions, 4 deletions

diff --git a/os.cpp b/os.cpp
index d437066..ae8caca 100644
--- a/os.cpp
+++ b/os.cpp
@@ -58,6 +58,25 @@ Os::Timestamp Os::GetTimestamp(clockid_t clock_id) const {
   return now_timestamp;
 }
 
+std::tuple<size_t, Os::Errno> Os::ReceiveDatagram(int fd, void* buf,
+                                                  size_t buflen) {
+  // recv() takes a size_t, but returns an ssize_t. That means that the largest
+  // successful read that recv() can report is the maximal ssize_t. Passing a
+  // larger |buflen| risks mistakenly reporting a truncated read.
+  CHECK(buflen <= GetMaxVal<ssize_t>());
+
+  const ssize_t res = raw_os_->Recv(fd, buf, buflen, MSG_TRUNC);
+  if (res < 0) {
+    return {0, errno};
+  }
+
+  // Due to the MSG_TRUNC flag, |res| may reasonably be larger than
+  // |buflen|. In such cases, |res| indicates the full size of the datagram,
+  // before being truncated to fit our buffer. Hence, we omit the
+  // buffer-overflow CHECK that exists in Write().
+  return {res, 0};
+}
+
 std::tuple<size_t, Os::Errno> Os::Write(int fd, const void* buf,
                                         size_t buflen) {
   // write() takes a size_t, but returns an ssize_t. That means that the
diff --git a/os.h b/os.h
index ebadd84..ca1e886 100644
--- a/os.h
+++ b/os.h
@@ -26,6 +26,7 @@
 
 #include "android-base/macros.h"
 
+#include "wifilogd/local_utils.h"
 #include "wifilogd/raw_os.h"
 
 namespace android {
@@ -60,6 +61,19 @@ class Os {
   // Returns the current time, as reported by the clock with |clock_id|.
   virtual Timestamp GetTimestamp(clockid_t clock_id) const;
 
+  // Receives a datagram of up to |buflen| from |fd|, writing the data to |buf|.
+  // Returns the size of the datagram, and the result of the operation (0 for
+  // success, |errno| otherwise).
+  //
+  // Notes:
+  // - |buflen| may not exceed the maximal value for ssize_t.
+  // - The call blocks until a datagram is available.
+  // - If the datagram is larger than |buflen|, only |buflen| bytes will
+  //   be received. The returned size_t will, however, reflect the full
+  //   length of the datagram.
+  virtual std::tuple<size_t, Errno> ReceiveDatagram(int fd, NONNULL void* buf,
+                                                    size_t buflen);
+
   // Writes |buflen| bytes from |buf| to |fd|. Returns the number of bytes
   // written, and the result of the operation (0 for success, |errno|
   // otherwise).
@@ -67,7 +81,7 @@ class Os {
   // Notes:
   // - |buflen| may not exceed the maximal value for ssize_t.
   // - The returned size_t will not exceed |buflen|.
-  virtual std::tuple<size_t, Errno> Write(int fd, const void* buf,
+  virtual std::tuple<size_t, Errno> Write(int fd, NONNULL const void* buf,
                                           size_t buflen);
 
  private:
diff --git a/raw_os.cpp b/raw_os.cpp
index 694f199..26aa1b8 100644
--- a/raw_os.cpp
+++ b/raw_os.cpp
@@ -33,6 +33,10 @@ int RawOs::ClockGettime(clockid_t clock_id, struct timespec* ts) const {
   return clock_gettime(clock_id, ts);
 }
 
+ssize_t RawOs::Recv(int sockfd, void* buf, size_t buflen, int flags) {
+  return recv(sockfd, buf, buflen, flags);
+}
+
 ssize_t RawOs::Write(int fd, const void* buf, size_t buflen) {
   return write(fd, buf, buflen);
 }
diff --git a/raw_os.h b/raw_os.h
index f5a17c4..6d74b6d 100644
--- a/raw_os.h
+++ b/raw_os.h
@@ -17,6 +17,8 @@
 #ifndef RAW_OS_H_
 #define RAW_OS_H_
 
+#include <sys/socket.h>
+#include <sys/types.h>
 #include <time.h>
 
 #include "android-base/macros.h"
@@ -36,6 +38,9 @@ class RawOs {
   virtual int ClockGettime(clockid_t clock_id,
                            NONNULL struct timespec* tspec) const;
 
+  // See recv().
+  virtual ssize_t Recv(int sockfd, void* buf, size_t buflen, int flags);
+
   // See write().
   virtual ssize_t Write(int fd, const void* buf, size_t buflen);
 
diff --git a/tests/mock_os.h b/tests/mock_os.h
index 354a6d1..5a378f3 100644
--- a/tests/mock_os.h
+++ b/tests/mock_os.h
@@ -33,6 +33,8 @@ class MockOs : public Os {
   ~MockOs() override;
 
   MOCK_CONST_METHOD1(GetTimestamp, Timestamp(clockid_t clock_id));
+  MOCK_METHOD3(ReceiveDatagram,
+               std::tuple<size_t, Errno>(int fd, void* buf, size_t buflen));
   MOCK_METHOD3(Write, std::tuple<size_t, Os::Errno>(int fd, const void* buf,
                                                     size_t buflen));
 
diff --git a/tests/mock_raw_os.h b/tests/mock_raw_os.h
index a39d085..13149d9 100644
--- a/tests/mock_raw_os.h
+++ b/tests/mock_raw_os.h
@@ -32,6 +32,7 @@ class MockRawOs : public RawOs {
 
   MOCK_CONST_METHOD2(ClockGettime,
                      int(clockid_t clock_id, struct timespec* tspec));
+  MOCK_METHOD4(Recv, ssize_t(int sockfd, void* buf, size_t buflen, int flags));
   MOCK_METHOD3(Write, ssize_t(int fd, const void* buf, size_t buflen));
 
  private:
diff --git a/tests/os_unittest.cpp b/tests/os_unittest.cpp
index 648a360..00dea33 100644
--- a/tests/os_unittest.cpp
+++ b/tests/os_unittest.cpp
@@ -65,6 +65,66 @@ TEST_F(OsTest, GetTimestampSucceeds) {
   EXPECT_EQ(kFakeNsecs, received.nsecs);
 }
 
+TEST_F(OsTest, ReceiveDatagramReturnsCorrectValueForMaxSizedDatagram) {
+  constexpr int kFakeFd = 100;
+  std::array<uint8_t, 8192> buffer{};
+  EXPECT_CALL(*raw_os_, Recv(kFakeFd, buffer.data(), buffer.size(), MSG_TRUNC))
+      .WillOnce(Return(buffer.size()));
+
+  constexpr std::tuple<size_t, Os::Errno> kExpectedResult{buffer.size(), 0};
+  EXPECT_EQ(kExpectedResult,
+            os_->ReceiveDatagram(kFakeFd, buffer.data(), buffer.size()));
+}
+
+TEST_F(OsTest, ReceieveDatagramReturnsCorrectValueForRegularSizedDatagram) {
+  constexpr int kFakeFd = 100;
+  constexpr auto kReadBufferSize = 8192;
+  constexpr auto kDatagramSize = kReadBufferSize / 2;
+  std::array<uint8_t, kReadBufferSize> buffer{};
+  EXPECT_CALL(*raw_os_, Recv(kFakeFd, buffer.data(), buffer.size(), MSG_TRUNC))
+      .WillOnce(Return(kDatagramSize));
+
+  constexpr std::tuple<size_t, Os::Errno> kExpectedResult{kDatagramSize, 0};
+  EXPECT_EQ(kExpectedResult,
+            os_->ReceiveDatagram(kFakeFd, buffer.data(), buffer.size()));
+}
+
+TEST_F(OsTest, ReceieveDatagramReturnsCorrectValueForOversizedDatagram) {
+  constexpr int kFakeFd = 100;
+  constexpr auto kReadBufferSize = 8192;
+  constexpr auto kDatagramSize = kReadBufferSize * 2;
+  std::array<uint8_t, kReadBufferSize> buffer{};
+  EXPECT_CALL(*raw_os_, Recv(kFakeFd, buffer.data(), buffer.size(), MSG_TRUNC))
+      .WillOnce(Return(kDatagramSize));
+
+  constexpr std::tuple<size_t, Os::Errno> kExpectedResult{kDatagramSize, 0};
+  EXPECT_EQ(kExpectedResult,
+            os_->ReceiveDatagram(kFakeFd, buffer.data(), buffer.size()));
+}
+
+TEST_F(OsTest, ReceieveDatagramReturnsCorrectValueForZeroByteDatagram) {
+  constexpr int kFakeFd = 100;
+  std::array<uint8_t, 8192> buffer{};
+  EXPECT_CALL(*raw_os_, Recv(kFakeFd, buffer.data(), buffer.size(), MSG_TRUNC))
+      .WillOnce(Return(0));
+
+  constexpr std::tuple<size_t, Os::Errno> kExpectedResult{0, 0};
+  EXPECT_EQ(kExpectedResult,
+            os_->ReceiveDatagram(kFakeFd, buffer.data(), buffer.size()));
+}
+
+TEST_F(OsTest, ReceieveDatagramReturnsCorrectValueOnFailure) {
+  constexpr int kFakeFd = 100;
+  constexpr Os::Errno kError = EBADF;
+  std::array<uint8_t, 8192> buffer{};
+  EXPECT_CALL(*raw_os_, Recv(kFakeFd, buffer.data(), buffer.size(), MSG_TRUNC))
+      .WillOnce(SetErrnoAndReturn(kError, -1));
+
+  constexpr std::tuple<size_t, Os::Errno> kExpectedResult{0, kError};
+  EXPECT_EQ(kExpectedResult,
+            os_->ReceiveDatagram(kFakeFd, buffer.data(), buffer.size()));
+}
+
 TEST_F(OsTest, WriteReturnsCorrectValueForSuccessfulWrite) {
   constexpr int kFakeFd = 100;
   constexpr std::array<uint8_t, 8192> buffer{};
@@ -86,9 +146,18 @@ TEST_F(OsTest, WriteReturnsCorrectValueForTruncatedWrite) {
   EXPECT_EQ(kExpectedResult, os_->Write(kFakeFd, buffer.data(), buffer.size()));
 }
 
+TEST_F(OsTest, WriteReturnsCorrectValueForSuccessfulZeroByteWrite) {
+  constexpr int kFakeFd = 100;
+  constexpr std::array<uint8_t, 0> buffer{};
+  EXPECT_CALL(*raw_os_, Write(kFakeFd, buffer.data(), 0)).WillOnce(Return(0));
+
+  constexpr std::tuple<size_t, Os::Errno> kExpectedResult{0, 0};
+  EXPECT_EQ(kExpectedResult, os_->Write(kFakeFd, buffer.data(), buffer.size()));
+}
+
 TEST_F(OsTest, WriteReturnsCorrectValueForFailedWrite) {
   constexpr int kFakeFd = 100;
-  constexpr Os::Errno kError = EBADFD;
+  constexpr Os::Errno kError = EBADF;
   constexpr std::array<uint8_t, 8192> buffer{};
   EXPECT_CALL(*raw_os_, Write(kFakeFd, buffer.data(), buffer.size()))
       .WillOnce(SetErrnoAndReturn(kError, -1));
@@ -97,6 +166,17 @@ TEST_F(OsTest, WriteReturnsCorrectValueForFailedWrite) {
   EXPECT_EQ(kExpectedResult, os_->Write(kFakeFd, buffer.data(), buffer.size()));
 }
 
+TEST_F(OsTest, WriteReturnsCorrectValueForFailedZeroByteWrite) {
+  constexpr int kFakeFd = 100;
+  constexpr Os::Errno kError = EBADF;
+  constexpr std::array<uint8_t, 0> buffer{};
+  EXPECT_CALL(*raw_os_, Write(kFakeFd, buffer.data(), 0))
+      .WillOnce(SetErrnoAndReturn(kError, -1));
+
+  constexpr std::tuple<size_t, Os::Errno> kExpectedResult{0, kError};
+  EXPECT_EQ(kExpectedResult, os_->Write(kFakeFd, buffer.data(), buffer.size()));
+}
+
 // Per
 // github.com/google/googletest/blob/master/googletest/docs/AdvancedGuide.md#death-tests,
 // death tests should be specially named.
@@ -116,11 +196,17 @@ TEST_F(OsDeathTest, GetTimestampRawOsErrorCausesDeath) {
   EXPECT_DEATH(os_->GetTimestamp(CLOCK_REALTIME), "Unexpected error");
 }
 
+TEST_F(OsDeathTest, ReceiveDatagramWithOverlyLargeBufferCausesDeath) {
+  constexpr int kFakeFd = 100;
+  std::array<uint8_t, 8192> buffer{};
+  EXPECT_DEATH(
+      os_->ReceiveDatagram(kFakeFd, buffer.data(), GetMaxVal<size_t>()),
+      "Check failed");
+}
+
 TEST_F(OsDeathTest, WriteWithOverlyLargeBufferCausesDeath) {
   constexpr int kFakeFd = 100;
   constexpr std::array<uint8_t, 8192> buffer{};
-  ON_CALL(*raw_os_, Write(kFakeFd, buffer.data(), GetMaxVal<size_t>()))
-      .WillByDefault(Return(GetMaxVal<ssize_t>()));
   EXPECT_DEATH(os_->Write(kFakeFd, buffer.data(), GetMaxVal<size_t>()),
                "Check failed");
 }