pw_transfer: Report errors from FinalizeWrite()

- Report errors from FinalizeWrite() in the final chunk of write
  transfers.
- Consolidate transfer completion logic in the FinishTransfer()
  function.
- Expand logging.

Change-Id: I994381c924e76545a6991d71c9e015c76c948ee6
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/62503
Pigweed-Auto-Submit: Wyatt Hepler <hepler@google.com>
Commit-Queue: Auto-Submit <auto-submit@pigweed.google.com.iam.gserviceaccount.com>
Reviewed-by: Keir Mierle <keir@google.com>

6 files changed, 99 insertions, 35 deletions

diff --git a/pw_transfer/docs.rst b/pw_transfer/docs.rst
index effad8c87..374d6486d 100644
--- a/pw_transfer/docs.rst
+++ b/pw_transfer/docs.rst
@@ -130,6 +130,11 @@ Errors
 At any point, either the client or server may terminate the transfer by sending
 an error chunk with the transfer ID and a non-OK status.
 
+- ``DATA_LOSS`` -- Failed to record data that was transferred. For example, a
+  flash write error occurred in Handler::FinalizeWrite().
+- ``INVALID_ARGUMENT`` -- Received a malformed chunk.
+- ``INTERNAL`` -- An assumption of the ``pw_transfer`` protocol was violated.
+
 Transmitter flow
 ================
 .. mermaid::
diff --git a/pw_transfer/public/pw_transfer/internal/server_context.h b/pw_transfer/public/pw_transfer/internal/server_context.h
index 05768e5d4..12fbe60a0 100644
--- a/pw_transfer/public/pw_transfer/internal/server_context.h
+++ b/pw_transfer/public/pw_transfer/internal/server_context.h
@@ -31,6 +31,8 @@ class ServerContext : public Context {
 
   constexpr bool active() const { return handler_ != nullptr; }
 
+  constexpr Type type() const { return type_; }
+
   // Begins a new transfer with the specified type and handler. Calls into the
   // handler's Prepare method.
   //
@@ -40,8 +42,10 @@ class ServerContext : public Context {
   // Ends the transfer with the given status, calling the handler's Finalize
   // method.
   //
+  // Returns DATA_LOSS if the finalize call fails.
+  //
   // Precondition: Transfer context is active.
-  void Finish(Status status);
+  Status Finish(Status status);
 
   stream::Reader& reader() const {
     PW_DASSERT(type_ == kRead);
diff --git a/pw_transfer/public/pw_transfer/transfer.h b/pw_transfer/public/pw_transfer/transfer.h
index 470cdf3d3..26a85a3b2 100644
--- a/pw_transfer/public/pw_transfer/transfer.h
+++ b/pw_transfer/public/pw_transfer/transfer.h
@@ -68,6 +68,9 @@ class TransferService : public generated::Transfer<TransferService> {
                        uint32_t transfer_id,
                        Status status);
 
+  // Calls transfer.Finish() and sends the final status chunk.
+  void FinishTransfer(internal::ServerContext& transfer, Status status);
+
   // Sends a out data chunk for a read transfer. Returns true if the data was
   // sent successfully.
   bool SendNextReadChunk(internal::ServerContext& context);
diff --git a/pw_transfer/server_context.cc b/pw_transfer/server_context.cc
index 4a7b18491..090d1f7e2 100644
--- a/pw_transfer/server_context.cc
+++ b/pw_transfer/server_context.cc
@@ -15,6 +15,7 @@
 #include "pw_transfer/internal/server_context.h"
 
 #include "pw_assert/check.h"
+#include "pw_log/log.h"
 #include "pw_status/try.h"
 #include "pw_varint/varint.h"
 
@@ -38,16 +39,26 @@ Status ServerContext::Start(Type type, Handler& handler) {
   return OkStatus();
 }
 
-void ServerContext::Finish(Status status) {
+Status ServerContext::Finish(const Status status) {
   PW_DCHECK(active());
 
+  Handler& handler = *handler_;
+  handler_ = nullptr;
+
   if (type_ == kRead) {
-    handler_->FinalizeRead(status);
-  } else {
-    handler_->FinalizeWrite(status);
+    handler.FinalizeRead(status);
+    return OkStatus();
   }
 
-  handler_ = nullptr;
+  if (Status finalized = handler.FinalizeWrite(status); !finalized.ok()) {
+    PW_LOG_ERROR(
+        "FinalizeWrite() for transfer %u failed with status %u; aborting with "
+        "DATA_LOSS",
+        static_cast<unsigned>(handler.id()),
+        static_cast<int>(finalized.code()));
+    return Status::DataLoss();
+  }
+  return OkStatus();
 }
 
 Result<ServerContext*> ServerContextPool::GetOrStartTransfer(uint32_t id) {
diff --git a/pw_transfer/transfer.cc b/pw_transfer/transfer.cc
index 7a1faf9da..86a9bb853 100644
--- a/pw_transfer/transfer.cc
+++ b/pw_transfer/transfer.cc
@@ -63,10 +63,8 @@ bool TransferService::SendNextReadChunk(internal::ServerContext& context) {
   // Begin by doing a partial encode of all the metadata fields, leaving the
   // buffer with usable space for the chunk data at the end.
   Chunk::MemoryEncoder encoder(buffer);
-  encoder.WriteTransferId(context.transfer_id())
-      .IgnoreError();  // TODO(pwbug/387): Handle Status properly
-  encoder.WriteOffset(context.offset())
-      .IgnoreError();  // TODO(pwbug/387): Handle Status properly
+  encoder.WriteTransferId(context.transfer_id()).IgnoreError();
+  encoder.WriteOffset(context.offset()).IgnoreError();
 
   // Reserve space for the data proto field overhead and use the remainder of
   // the buffer for the chunk data.
@@ -83,22 +81,34 @@ bool TransferService::SendNextReadChunk(internal::ServerContext& context) {
   Result<ByteSpan> data = context.reader().Read(data_buffer);
   if (data.status().IsOutOfRange()) {
     // No more data to read.
-    encoder.WriteRemainingBytes(0)
-        .IgnoreError();  // TODO(pwbug/387): Handle Status properly
+    encoder.WriteRemainingBytes(0).IgnoreError();
     context.set_pending_bytes(0);
-  } else if (!data.ok()) {
-    read_stream_.ReleaseBuffer();
-    return false;
-  } else {
-    encoder.WriteData(data.value())
-        .IgnoreError();  // TODO(pwbug/387): Handle Status properly
+  } else if (data.ok()) {
+    encoder.WriteData(data.value()).IgnoreError();
     context.set_offset(context.offset() + data.value().size());
     context.set_pending_bytes(context.pending_bytes() - data.value().size());
+  } else {
+    read_stream_.ReleaseBuffer();
+    return false;
   }
 
   return read_stream_.Write(encoder).ok();
 }
 
+void TransferService::FinishTransfer(internal::ServerContext& transfer,
+                                     Status status) {
+  const uint32_t id = transfer.transfer_id();
+  PW_LOG_INFO("Sending status %u in final chunk of transfer %u",
+              status.code(),
+              static_cast<unsigned>(id));
+  status.Update(transfer.Finish(status));
+
+  RawServerReaderWriter& stream =
+      transfer.type() == internal::ServerContext::kRead ? read_stream_
+                                                        : write_stream_;
+  SendStatusChunk(stream, id, status);
+}
+
 void TransferService::OnReadMessage(ConstByteSpan message) {
   // All incoming chunks in a client read transfer are transfer parameter
   // updates, except for the final chunk, which is an acknowledgement of
@@ -138,19 +148,17 @@ void TransferService::OnReadMessage(ConstByteSpan message) {
     // Transfer has been terminated (successfully or not).
     Status status = parameters.status.value();
     if (!status.ok()) {
-      PW_LOG_ERROR("Transfer %u failed with status %d",
+      PW_LOG_ERROR("Receiver terminated read transfer %u with status %d",
                    static_cast<unsigned>(parameters.transfer_id),
                    static_cast<int>(status.code()));
     }
-    transfer.Finish(status);
+    transfer.Finish(status).IgnoreError();
     return;
   }
 
   if (!parameters.pending_bytes.has_value()) {
     // Malformed chunk.
-    SendStatusChunk(
-        read_stream_, parameters.transfer_id, Status::InvalidArgument());
-    transfer.Finish(Status::InvalidArgument());
+    FinishTransfer(transfer, Status::InvalidArgument());
     return;
   }
 
@@ -162,9 +170,7 @@ void TransferService::OnReadMessage(ConstByteSpan message) {
     //   transfer.set_offset(parameters.offset.value());
     //   transfer.Seek(transfer.offset());
     //
-    SendStatusChunk(
-        read_stream_, parameters.transfer_id, Status::Unimplemented());
-    transfer.Finish(Status::Unimplemented());
+    FinishTransfer(transfer, Status::Unimplemented());
     return;
   }
 
@@ -207,7 +213,7 @@ void TransferService::OnWriteMessage(ConstByteSpan message) {
 
   // Check for a client-side error terminating the transfer.
   if (chunk.status.has_value()) {
-    transfer.Finish(chunk.status.value());
+    transfer.Finish(chunk.status.value()).IgnoreError();
     return;
   }
 
@@ -222,8 +228,13 @@ void TransferService::OnWriteMessage(ConstByteSpan message) {
       chunk_data_processed = true;
     } else if (chunk.data.size() <= transfer.pending_bytes()) {
       if (Status status = transfer.writer().Write(chunk.data); !status.ok()) {
-        SendStatusChunk(write_stream_, chunk.transfer_id, status);
-        transfer.Finish(status);
+        PW_LOG_ERROR(
+            "Transfer %u write of %u B chunk failed with status %u; aborting "
+            "with DATA_LOSS",
+            static_cast<unsigned>(chunk.transfer_id),
+            static_cast<unsigned>(chunk.data.size()),
+            status.code());
+        FinishTransfer(transfer, Status::DataLoss());
         return;
       }
       transfer.set_offset(transfer.offset() + chunk.data.size());
@@ -235,8 +246,7 @@ void TransferService::OnWriteMessage(ConstByteSpan message) {
       // could potentially result in an infinite transfer loop.
       PW_LOG_ERROR(
           "Received more data than what was requested; terminating transfer.");
-      SendStatusChunk(write_stream_, chunk.transfer_id, Status::Internal());
-      transfer.Finish(Status::Internal());
+      FinishTransfer(transfer, Status::Internal());
       return;
     }
   } else {
@@ -247,8 +257,7 @@ void TransferService::OnWriteMessage(ConstByteSpan message) {
   // When the client sets remaining_bytes to 0, it indicates completion of the
   // transfer. Acknowledge the completion through a status chunk and clean up.
   if (chunk_data_processed && chunk.remaining_bytes == 0) {
-    SendStatusChunk(write_stream_, chunk.transfer_id, OkStatus());
-    transfer.Finish(OkStatus());
+    FinishTransfer(transfer, OkStatus());
     return;
   }
 
@@ -287,7 +296,7 @@ void TransferService::OnWriteMessage(ConstByteSpan message) {
                  static_cast<unsigned>(parameters.transfer_id),
                  data.status().code());
     write_stream_.ReleaseBuffer();
-    transfer.Finish(Status::Internal());
+    FinishTransfer(transfer, Status::Internal());
   }
 }
 
diff --git a/pw_transfer/transfer_test.cc b/pw_transfer/transfer_test.cc
index c1e6fd6d4..046936ea0 100644
--- a/pw_transfer/transfer_test.cc
+++ b/pw_transfer/transfer_test.cc
@@ -351,7 +351,11 @@ class SimpleWriteTransfer final : public WriteOnlyHandler {
   Status FinalizeWrite(Status status) final {
     finalize_write_called = true;
     finalize_write_status = status;
-    return OkStatus();
+    return finalize_write_return_status_;
+  }
+
+  void set_finalize_write_return(Status status) {
+    finalize_write_return_status_ = status;
   }
 
   bool prepare_write_called;
@@ -359,6 +363,7 @@ class SimpleWriteTransfer final : public WriteOnlyHandler {
   Status finalize_write_status;
 
  private:
+  Status finalize_write_return_status_;
   stream::MemoryWriter writer_;
 };
 
@@ -402,6 +407,33 @@ TEST(Transfer, Write_SingleChunk) {
   EXPECT_EQ(std::memcmp(buffer.data(), data.data(), data.size()), 0);
 }
 
+TEST(Transfer, Write_FinalizeFails) {
+  constexpr auto data = bytes::Initialized<32>([](size_t i) { return i; });
+  std::array<std::byte, sizeof(data)> buffer = {};
+  SimpleWriteTransfer handler(7, buffer);
+
+  // Return an error when FinalizeWrite is called.
+  handler.set_finalize_write_return(Status::FailedPrecondition());
+
+  PW_RAW_TEST_METHOD_CONTEXT(TransferService, Write) ctx(64, 64);
+  ctx.service().RegisterHandler(handler);
+
+  ctx.call();
+  ctx.SendClientStream(EncodeChunk({.transfer_id = 7}));
+  ctx.SendClientStream<64>(EncodeChunk({.transfer_id = 7,
+                                        .offset = 0,
+                                        .data = std::span(data),
+                                        .remaining_bytes = 0}));
+
+  Chunk chunk = DecodeChunk(ctx.responses()[1]);
+  EXPECT_EQ(chunk.transfer_id, 7u);
+  ASSERT_TRUE(chunk.status.has_value());
+  EXPECT_EQ(chunk.status.value(), Status::DataLoss());
+
+  EXPECT_TRUE(handler.finalize_write_called);
+  EXPECT_EQ(handler.finalize_write_status, OkStatus());
+}
+
 TEST(Transfer, Write_MultiChunk) {
   constexpr auto data = bytes::Initialized<32>([](size_t i) { return i; });
   std::array<std::byte, sizeof(data)> buffer = {};