From a2bdd10fccb346d02f12d3961d6a5b2d281c5d56 Mon Sep 17 00:00:00 2001
From: Duong Nguyen <duongnt.is@gmail.com>
Date: Mon, 3 Jun 2024 10:35:07 -0700
Subject: [PATCH] RATIS-2094. Avoid corruptions from TransactionContext's
 stateMachineLogEntry and stateMachineContext. (#1106)

---
 .../org/apache/ratis/util/ReferenceCountedObject.java |  8 ++++++++
 .../examples/filestore/FileStoreStateMachine.java     | 11 ++++++-----
 .../apache/ratis/statemachine/TransactionContext.java |  3 +++
 .../apache/ratis/statemachine/TestStateMachine.java   |  2 +-
 4 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/ratis-common/src/main/java/org/apache/ratis/util/ReferenceCountedObject.java b/ratis-common/src/main/java/org/apache/ratis/util/ReferenceCountedObject.java
index eb5ff30ae8..b2c53182d3 100644
--- a/ratis-common/src/main/java/org/apache/ratis/util/ReferenceCountedObject.java
+++ b/ratis-common/src/main/java/org/apache/ratis/util/ReferenceCountedObject.java
@@ -23,6 +23,7 @@
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 /**
  * A reference-counted object can be retained for later use
@@ -152,6 +153,13 @@ public boolean release() {
     };
   }
 
+  /**
+   * @return a {@link ReferenceCountedObject} by apply the given function to this object.
+   */
+  default <V> ReferenceCountedObject<V> apply(Function<T, V> function) {
+    return delegate(function.apply(get()));
+  }
+
   /**
    * Wrap the given value as a {@link ReferenceCountedObject}.
    *
diff --git a/ratis-examples/src/main/java/org/apache/ratis/examples/filestore/FileStoreStateMachine.java b/ratis-examples/src/main/java/org/apache/ratis/examples/filestore/FileStoreStateMachine.java
index f870cbacfe..2fb9cace66 100644
--- a/ratis-examples/src/main/java/org/apache/ratis/examples/filestore/FileStoreStateMachine.java
+++ b/ratis-examples/src/main/java/org/apache/ratis/examples/filestore/FileStoreStateMachine.java
@@ -115,11 +115,12 @@ public TransactionContext startTransaction(RaftClientRequest request) throws IOE
 
   @Override
   public TransactionContext startTransaction(LogEntryProto entry, RaftProtos.RaftPeerRole role) {
+    ByteString copied = ByteString.copyFrom(entry.getStateMachineLogEntry().getLogData().asReadOnlyByteBuffer());
     return TransactionContext.newBuilder()
         .setStateMachine(this)
         .setLogEntry(entry)
         .setServerRole(role)
-        .setStateMachineContext(getProto(entry))
+        .setStateMachineContext(getProto(copied))
         .build();
   }
 
@@ -147,14 +148,14 @@ static FileStoreRequestProto getProto(TransactionContext context, LogEntryProto
         return proto;
       }
     }
-    return getProto(entry);
+    return getProto(entry.getStateMachineLogEntry().getLogData());
   }
 
-  static FileStoreRequestProto getProto(LogEntryProto entry) {
+  static FileStoreRequestProto getProto(ByteString bytes) {
     try {
-      return FileStoreRequestProto.parseFrom(entry.getStateMachineLogEntry().getLogData());
+      return FileStoreRequestProto.parseFrom(bytes);
     } catch (InvalidProtocolBufferException e) {
-      throw new IllegalArgumentException("Failed to parse data, entry=" + entry, e);
+      throw new IllegalArgumentException("Failed to parse data", e);
     }
   }
 
diff --git a/ratis-server-api/src/main/java/org/apache/ratis/statemachine/TransactionContext.java b/ratis-server-api/src/main/java/org/apache/ratis/statemachine/TransactionContext.java
index 2ec87e37ad..6f08c4a55d 100644
--- a/ratis-server-api/src/main/java/org/apache/ratis/statemachine/TransactionContext.java
+++ b/ratis-server-api/src/main/java/org/apache/ratis/statemachine/TransactionContext.java
@@ -60,7 +60,10 @@ public interface TransactionContext {
   /**
    * Returns the data from the {@link StateMachine}
    * @return the data from the {@link StateMachine}
+   * @deprecated access StateMachineLogEntry via {@link TransactionContext#getLogEntryRef()} or
+   * {@link TransactionContext#getLogEntryUnsafe()}
    */
+  @Deprecated
   StateMachineLogEntryProto getStateMachineLogEntry();
 
   /** Set exception in case of failure. */
diff --git a/ratis-test/src/test/java/org/apache/ratis/statemachine/TestStateMachine.java b/ratis-test/src/test/java/org/apache/ratis/statemachine/TestStateMachine.java
index 07ea4edbcb..0c49b31af2 100644
--- a/ratis-test/src/test/java/org/apache/ratis/statemachine/TestStateMachine.java
+++ b/ratis-test/src/test/java/org/apache/ratis/statemachine/TestStateMachine.java
@@ -92,7 +92,7 @@ public TransactionContext startTransaction(RaftClientRequest request) {
     public CompletableFuture<Message> applyTransaction(TransactionContext trx) {
       try {
         assertNotNull(trx.getLogEntryUnsafe());
-        assertNotNull(trx.getStateMachineLogEntry());
+        assertNotNull(trx.getLogEntryUnsafe().getStateMachineLogEntry());
         Object context = trx.getStateMachineContext();
         if (isLeader.get()) {
           assertNotNull(trx.getClientRequest());