From 1a78b33a897972ddadccdc5569ff00c53df55902 Mon Sep 17 00:00:00 2001 From: Janusz Sobczak Date: Wed, 22 Mar 2023 17:44:00 -0700 Subject: [PATCH] Add InputStream::ReadExactly Read() can return fewer bytes than requested. ReadExactly() will call Read() repeatedly until we have read as many bytes as we need. PiperOrigin-RevId: 518719418 --- Package.swift | 1 + .../implementation/base_endpoint_channel.cc | 27 +--- fastpair/message_stream/medium.cc | 4 +- internal/platform/BUILD | 1 + internal/platform/exception.h | 9 ++ internal/platform/input_stream.cc | 31 +++++ internal/platform/input_stream.h | 5 + internal/platform/input_stream_test.cc | 119 ++++++++++++++++++ 8 files changed, 170 insertions(+), 27 deletions(-) create mode 100644 internal/platform/input_stream_test.cc diff --git a/Package.swift b/Package.swift index d1c82fc8..dfdaaa94 100644 --- a/Package.swift +++ b/Package.swift @@ -493,6 +493,7 @@ let package = Package( "internal/platform/byte_array_test.cc", "internal/platform/bluetooth_utils_test.cc", "internal/platform/credential_storage_impl_test.cc", + "internal/platform/input_stream_test.cc", "internal/platform/single_thread_executor_test.cc", "internal/platform/scheduled_executor_test.cc", "internal/platform/count_down_latch_test.cc", diff --git a/connections/implementation/base_endpoint_channel.cc b/connections/implementation/base_endpoint_channel.cc index 4aaf3589..b2b2f96e 100644 --- a/connections/implementation/base_endpoint_channel.cc +++ b/connections/implementation/base_endpoint_channel.cc @@ -54,31 +54,8 @@ ByteArray IntToBytes(std::int32_t value) { return ByteArray(int_bytes, sizeof(int_bytes)); } -ExceptionOr ReadExactly(InputStream* reader, std::int64_t size) { - ByteArray buffer(size); - std::int64_t current_pos = 0; - - while (current_pos < size) { - ExceptionOr read_bytes = reader->Read(size - current_pos); - if (!read_bytes.ok()) { - return read_bytes; - } - ByteArray result = read_bytes.result(); - - if (result.Empty()) { - NEARBY_LOGS(WARNING) << __func__ << ": Empty result when reading bytes."; - return ExceptionOr(Exception::kIo); - } - - buffer.CopyAt(current_pos, result); - current_pos += result.size(); - } - - return ExceptionOr(std::move(buffer)); -} - ExceptionOr ReadInt(InputStream* reader) { - ExceptionOr read_bytes = ReadExactly(reader, sizeof(std::int32_t)); + ExceptionOr read_bytes = reader->ReadExactly(sizeof(std::int32_t)); if (!read_bytes.ok()) { return ExceptionOr(read_bytes.exception()); } @@ -147,7 +124,7 @@ ExceptionOr BaseEndpointChannel::Read( return ExceptionOr(Exception::kIo); } - ExceptionOr read_bytes = ReadExactly(reader_, read_int.result()); + ExceptionOr read_bytes = reader_->ReadExactly(read_int.result()); if (!read_bytes.ok()) { return read_bytes; } diff --git a/fastpair/message_stream/medium.cc b/fastpair/message_stream/medium.cc index 5e8bd4d7..0ceaa3f7 100644 --- a/fastpair/message_stream/medium.cc +++ b/fastpair/message_stream/medium.cc @@ -102,7 +102,7 @@ void Medium::RunLoop(BluetoothSocket socket) { NEARBY_LOGS(INFO) << "Run loop"; InputStream& input = socket.GetInputStream(); while (!cancellation_flag_.Cancelled()) { - ExceptionOr header = input.Read(kHeaderSize); + ExceptionOr header = input.ReadExactly(kHeaderSize); if (!header.ok() || header.result().size() != kHeaderSize) { break; } @@ -113,7 +113,7 @@ void Medium::RunLoop(BluetoothSocket socket) { static_cast(data[3]); ExceptionOr payload; if (length > 0) { - payload = input.Read(length); + payload = input.ReadExactly(length); } else { payload = ExceptionOr(ByteArray("")); } diff --git a/internal/platform/BUILD b/internal/platform/BUILD index 455e1d56..8fea2d79 100644 --- a/internal/platform/BUILD +++ b/internal/platform/BUILD @@ -232,6 +232,7 @@ cc_test( "bluetooth_utils_test.cc", "byte_array_test.cc", "feature_flags_test.cc", + "input_stream_test.cc", "prng_test.cc", ], deps = [ diff --git a/internal/platform/exception.h b/internal/platform/exception.h index 2ce4618b..a151a103 100644 --- a/internal/platform/exception.h +++ b/internal/platform/exception.h @@ -107,6 +107,15 @@ class ExceptionOr { Exception exception_{Exception::kFailed}; }; +template +constexpr inline bool operator==(const ExceptionOr& a, + const ExceptionOr& b) { + if (a.ok() && b.ok()) { + return a.result() == b.result(); + } + return a.exception() == b.exception(); +} + } // namespace nearby #endif // PLATFORM_BASE_EXCEPTION_H_ diff --git a/internal/platform/input_stream.cc b/internal/platform/input_stream.cc index d9043e7e..0d18cfa4 100644 --- a/internal/platform/input_stream.cc +++ b/internal/platform/input_stream.cc @@ -17,6 +17,7 @@ #include #include #include +#include #include "internal/platform/byte_array.h" #include "internal/platform/exception.h" @@ -43,4 +44,34 @@ ExceptionOr InputStream::Skip(size_t offset) { return ExceptionOr(offset); } +ExceptionOr InputStream::ReadExactly(std::size_t size) { + ByteArray buffer; + std::size_t current_pos = 0; + + while (current_pos < size) { + ExceptionOr read_bytes = Read(size - current_pos); + if (!read_bytes.ok()) { + return read_bytes; + } + const ByteArray& result = read_bytes.result(); + + if (result.Empty()) { + return ExceptionOr(Exception::kIo); + } + if (current_pos == 0) { + if (result.size() == size) { + // We have read the requested `size` bytes in one chunk. We can return + // it directly. + return read_bytes; + } else { + // Reserve space for in the buffer. + buffer.SetData(size); + } + } + buffer.CopyAt(current_pos, result); + current_pos += result.size(); + } + + return ExceptionOr(std::move(buffer)); +} } // namespace nearby diff --git a/internal/platform/input_stream.h b/internal/platform/input_stream.h index f0315300..383964b6 100644 --- a/internal/platform/input_stream.h +++ b/internal/platform/input_stream.h @@ -38,6 +38,11 @@ class InputStream { // or Exception::kIo on error. virtual ExceptionOr Skip(size_t offset); + // Reads exactly `size` bytes from the input stream. + // Return Exception::kIo on error, or if end of file is reached before reading + // `size` bytes. + ExceptionOr ReadExactly(std::size_t size); + // throws Exception::kIo virtual Exception Close() = 0; }; diff --git a/internal/platform/input_stream_test.cc b/internal/platform/input_stream_test.cc new file mode 100644 index 00000000..53d9d4c0 --- /dev/null +++ b/internal/platform/input_stream_test.cc @@ -0,0 +1,119 @@ +// Copyright 2023 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "internal/platform/input_stream.h" + +#include + +#include "gmock/gmock.h" +#include "protobuf-matchers/protocol-buffer-matchers.h" +#include "gtest/gtest.h" +#include "internal/platform/byte_array.h" +#include "internal/platform/exception.h" + +namespace nearby { +namespace { + +using ::testing::_; +using ::testing::InSequence; +using ::testing::NiceMock; +using ::testing::Return; + +class TestInputStream : public InputStream { + public: + MOCK_METHOD(ExceptionOr, Read, (std::int64_t), (override)); + MOCK_METHOD(Exception, Close, (), (override)); +}; + +// Returns a ByteArray with values: a, a + 1, a + 2, ..., b - 1. +ExceptionOr Range(char a, char b) { + std::string s; + for (char c = a; c < b; c++) { + s.push_back(c); + } + return ExceptionOr(ByteArray(s)); +} + +TEST(InputStreamTest, Skip) { + NiceMock stream; + InSequence seq; + EXPECT_CALL(stream, Read(50)).WillOnce(Return(Range(0, 10))); + EXPECT_CALL(stream, Read(40)).WillOnce(Return(Range(0, 10))); + EXPECT_CALL(stream, Read(30)).WillOnce(Return(Range(0, 30))); + + ExceptionOr skipped = stream.Skip(50); + + EXPECT_EQ(skipped.result(), 50); +} + +TEST(InputStreamTest, SkipsLessOnEof) { + NiceMock stream; + InSequence seq; + EXPECT_CALL(stream, Read(50)).WillOnce(Return(Range(0, 10))); + EXPECT_CALL(stream, Read(40)).WillOnce(Return(Range(0, 10))); + // Returns EOF + EXPECT_CALL(stream, Read(30)).WillOnce(Return(Range(0, 0))); + + ExceptionOr skipped = stream.Skip(50); + + EXPECT_EQ(skipped.result(), 20); +} + +TEST(InputStreamTest, SkipFailsOnError) { + NiceMock stream; + InSequence seq; + EXPECT_CALL(stream, Read(50)).WillOnce(Return(Range(0, 10))); + EXPECT_CALL(stream, Read(40)) + .WillOnce(Return(ExceptionOr(Exception::kIo))); + + ExceptionOr skipped = stream.Skip(50); + + EXPECT_EQ(skipped.exception(), Exception::kIo); +} + +TEST(InputStreamTest, ReadExactlyOneChunk) { + NiceMock stream; + EXPECT_CALL(stream, Read(10)).WillOnce(Return(Range(0, 10))); + + ExceptionOr result = stream.ReadExactly(10); + + EXPECT_EQ(result, Range(0, 10)); +} + +TEST(InputStreamTest, ReadExactlyMultipleChunks) { + NiceMock stream; + InSequence seq; + EXPECT_CALL(stream, Read(30)).WillOnce(Return(Range(0, 10))); + EXPECT_CALL(stream, Read(20)).WillOnce(Return(Range(10, 20))); + EXPECT_CALL(stream, Read(10)).WillOnce(Return(Range(20, 30))); + + ExceptionOr result = stream.ReadExactly(30); + + EXPECT_EQ(result, Range(0, 30)); +} + +TEST(InputStreamTest, ReadExactlyFailsOnError) { + NiceMock stream; + InSequence seq; + EXPECT_CALL(stream, Read(30)).WillOnce(Return(Range(0, 10))); + EXPECT_CALL(stream, Read(20)) + .WillOnce(Return(ExceptionOr(Exception::kIo))); + + ExceptionOr result = stream.ReadExactly(30); + + EXPECT_EQ(result.exception(), Exception::kIo); +} + +} // namespace +} // namespace nearby