ampr: fix APR completion loss in 768-byte read batches

Co-authored-by: Chai T <ct@talnikar.in>

src/libs/libAmpr.cpp

@@ -868,6 +868,10 @@ static uint64_t AprReadFileRecordSize(uint64_t file_offset) {
 	return (file_offset >> 32u) != 0 ? READ_FILE_RECORD_SIZE_EXT : READ_FILE_RECORD_SIZE;
 }
 
+static uint64_t WriteAddressRecordSize(uint64_t value) {
+	return value <= 3 ? 8 : value < (uint64_t {1} << 34) ? 12 : 16;
+}
+
 static uint64_t AprReadGatherRecordSize(uint64_t file_offset) {
 	return file_offset > 0x3ffffu ? READ_GATHER_RECORD_SIZE_EXT : READ_GATHER_RECORD_SIZE;
 }
@@ -1617,11 +1621,10 @@ static uint64_t KYTY_SYSV_ABI MeasureCommandSizeFixed32(uint64_t, uint64_t, uint
 	return 0x20;
 }
 
-static uint64_t KYTY_SYSV_ABI MeasureCommandSizeWriteAddressOnCompletion(volatile uint64_t*,
+static uint64_t KYTY_SYSV_ABI MeasureCommandSizeWriteAddress(volatile uint64_t*, uint64_t value) {
-                                                                         uint64_t) {
 	PRINT_NAME();
 
-	return 0x20;
+	return WriteAddressRecordSize(value);
 }
 
 static uint64_t KYTY_SYSV_ABI
@@ -1950,7 +1953,7 @@ static int KYTY_SYSV_ABI CommandBufferWaitOnCounter(void* command_buffer, uint8_
 }
 
 static int AppendWriteAddressCommand(void* command_buffer, volatile uint64_t* address,
-                                     uint64_t value) {
+                                     uint64_t value, uint64_t record_size) {
 	if (command_buffer == nullptr || address == nullptr ||
 	    (reinterpret_cast<uint64_t>(address) & 7) != 0) {
 		return LibKernel::KERNEL_ERROR_EINVAL;
@@ -1958,7 +1961,7 @@ static int AppendWriteAddressCommand(void* command_buffer, volatile uint64_t* ad
 	return AppendCommandRecord(
 	           reinterpret_cast<uint64_t>(command_buffer),
 	           CommandBufferState::WriteAddressCommand {reinterpret_cast<uint64_t>(address), value},
-	           0x20)
+	           record_size)
 	           ? OK
 	           : LibKernel::KERNEL_ERROR_EBUSY;
 }
@@ -1967,7 +1970,7 @@ static int KYTY_SYSV_ABI CommandBufferWriteAddress(void* command_buffer, volatil
                                                    uint64_t value, uint32_t) {
 	PRINT_NAME();
 
-	return AppendWriteAddressCommand(command_buffer, address, value);
+	return AppendWriteAddressCommand(command_buffer, address, value, WriteAddressRecordSize(value));
 }
 
 static int KYTY_SYSV_ABI CommandBufferWriteAddressOnCompletion(void*              command_buffer,
@@ -1975,7 +1978,7 @@ static int KYTY_SYSV_ABI CommandBufferWriteAddressOnCompletion(void*
                                                                uint64_t           value) {
 	PRINT_NAME();
 
-	return AppendWriteAddressCommand(command_buffer, address, value);
+	return AppendWriteAddressCommand(command_buffer, address, value, WriteAddressRecordSize(value));
 }
 
 static int KYTY_SYSV_ABI CommandBufferWriteCounter(void* command_buffer, uint8_t index,
@@ -2007,14 +2010,14 @@ static int KYTY_SYSV_ABI CommandBufferWriteAddressFromTimeCounter(void*
                                                                   uint32_t) {
 	PRINT_NAME();
 
-	return AppendWriteAddressCommand(command_buffer, address, uint64_t {0});
+	return AppendWriteAddressCommand(command_buffer, address, uint64_t {0}, 0x20);
 }
 
 static int KYTY_SYSV_ABI CommandBufferWriteAddressFromTimeCounterOnCompletion(
     void* command_buffer, volatile uint64_t* address) {
 	PRINT_NAME();
 
-	return AppendWriteAddressCommand(command_buffer, address, uint64_t {0});
+	return AppendWriteAddressCommand(command_buffer, address, uint64_t {0}, 0x20);
 }
 
 static int KYTY_SYSV_ABI CommandBufferWriteAddressFromCounter(void*              command_buffer,
@@ -2022,14 +2025,14 @@ static int KYTY_SYSV_ABI CommandBufferWriteAddressFromCounter(void*
                                                               uint32_t) {
 	PRINT_NAME();
 
-	return AppendWriteAddressCommand(command_buffer, address, uint64_t {0});
+	return AppendWriteAddressCommand(command_buffer, address, uint64_t {0}, 0x20);
 }
 
 static int KYTY_SYSV_ABI CommandBufferWriteAddressFromCounterOnCompletion(
     void* command_buffer, volatile uint64_t* address, uint8_t) {
 	PRINT_NAME();
 
-	return AppendWriteAddressCommand(command_buffer, address, uint64_t {0});
+	return AppendWriteAddressCommand(command_buffer, address, uint64_t {0}, 0x20);
 }
 
 static int KYTY_SYSV_ABI CommandBufferWriteAddressFromCounterPair(void*              command_buffer,
@@ -2037,14 +2040,14 @@ static int KYTY_SYSV_ABI CommandBufferWriteAddressFromCounterPair(void*
                                                                   uint8_t, uint32_t) {
 	PRINT_NAME();
 
-	return AppendWriteAddressCommand(command_buffer, address, uint64_t {0});
+	return AppendWriteAddressCommand(command_buffer, address, uint64_t {0}, 0x20);
 }
 
 static int KYTY_SYSV_ABI CommandBufferWriteAddressFromCounterPairOnCompletion(
     void* command_buffer, volatile uint64_t* address, uint8_t) {
 	PRINT_NAME();
 
-	return AppendWriteAddressCommand(command_buffer, address, uint64_t {0});
+	return AppendWriteAddressCommand(command_buffer, address, uint64_t {0}, 0x20);
 }
 
 static int KYTY_SYSV_ABI AprCommandBufferReadFileGather(void*    command_buffer, uint64_t, uint64_t,
@@ -2452,14 +2455,14 @@ LIB_DEFINE(InitAmpr_1) {
 	LIB_FUNC("vWU-odnS+fU", Ampr::MeasureCommandSizeReadFile);
 	LIB_FUNC("sSAUCCU1dv4", Ampr::MeasureCommandSizeWriteKernelEventQueue);
 	LIB_FUNC("Zi3dBUjgyXI", Ampr::MeasureCommandSizeWriteKernelEventQueue);
-	LIB_FUNC("C+IEj+BsAFM", Ampr::MeasureCommandSizeWriteAddressOnCompletion);
+	LIB_FUNC("C+IEj+BsAFM", Ampr::MeasureCommandSizeWriteAddress);
 	LIB_FUNC("x7SQEXfeovg", Ampr::MeasureCommandSizeWriteAddressFromTimeCounterOnCompletion);
 	LIB_FUNC("32AcaTaBPSY", Ampr::MeasureCommandSizeWriteAddressFromCounterOnCompletion);
 	LIB_FUNC("vxC58+DRk-U", Ampr::MeasureCommandSizeWriteAddressFromCounterPairOnCompletion);
 	LIB_FUNC("4muPEJ-x5N8", Ampr::MeasureCommandSizeWriteCounterOnCompletion);
 	LIB_FUNC("0BMj1hgG+kE", Ampr::MeasureCommandSizeFixed32);
 	LIB_FUNC("ClnsFLLLcss", Ampr::MeasureCommandSizeFixed32);
-	LIB_FUNC("4fgtGfXDrFc", Ampr::MeasureCommandSizeFixed32);
+	LIB_FUNC("4fgtGfXDrFc", Ampr::MeasureCommandSizeWriteAddress);
 	LIB_FUNC("gAtc79UTt5E", Ampr::MeasureCommandSizeFixed32);
 	LIB_FUNC("JYd9g9L+TmE", Ampr::MeasureCommandSizeFixed32);
 	LIB_FUNC("2Hw8gjMdwSY", Ampr::MeasureCommandSizeFixed32);

tests/KernelFileSystemTests.cpp

@@ -898,6 +898,11 @@ void CheckAmprOrdering(Loader::SymbolDatabase &symbols, uint32_t file_id) {
   using WriteCounter = int (KYTY_SYSV_ABI *)(void *, uint8_t, uint8_t, uint64_t,
                                            uint8_t, uint32_t);
   using WriteAddress = int (KYTY_SYSV_ABI *)(void *, volatile uint64_t *, uint64_t);
+  using WriteImmediately = int (KYTY_SYSV_ABI *)(void *, volatile uint64_t *, uint64_t,
+                                               uint32_t);
+  using MeasureWriteAddress = uint64_t (KYTY_SYSV_ABI *)(volatile uint64_t *, uint64_t);
+  using MeasureWriteImmediately = uint64_t (KYTY_SYSV_ABI *)(volatile uint64_t *, uint64_t,
+                                                           uint32_t);
   using ReadFile = int (KYTY_SYSV_ABI *)(void *, uint64_t, uint64_t, uint32_t,
                                        void *, uint64_t, uint64_t);
   using WriteKernelEvent = int (KYTY_SYSV_ABI *)(void *, uint64_t, uint64_t, uint64_t,
@@ -916,6 +921,9 @@ void CheckAmprOrdering(Loader::SymbolDatabase &symbols, uint32_t file_id) {
   const auto wait_counter = reinterpret_cast<WaitCounter>(find("cQb8Zr8Q0Y0"));
   const auto write_counter = reinterpret_cast<WriteCounter>(find("jK+yuYCI7MA"));
   const auto write_address = reinterpret_cast<WriteAddress>(find("sJXyWHjP-F8"));
+  const auto write_immediately = reinterpret_cast<WriteImmediately>(find("j0+3uJMxYJY"));
+  const auto measure_write_address = reinterpret_cast<MeasureWriteAddress>(find("C+IEj+BsAFM"));
+  const auto measure_write_immediately = reinterpret_cast<MeasureWriteImmediately>(find("4fgtGfXDrFc"));
   const auto read_file = reinterpret_cast<ReadFile>(find("mQ16-QdKv7k"));
   const auto write_event = reinterpret_cast<WriteKernelEvent>(find("H896Pt-yB4I"));
   const auto submit_apr = reinterpret_cast<SubmitApr>(find("ASoW5WE-UPo"));
@@ -927,7 +935,7 @@ void CheckAmprOrdering(Loader::SymbolDatabase &symbols, uint32_t file_id) {
     std::array<uint32_t, 256> data {};
   };
   std::array<Buffer, 4> buffers;
-  const auto reset = [&](size_t apr_count) {
+  const auto reset = [&](size_t apr_count, uint32_t data_size = sizeof(Buffer::data)) {
     for (size_t i = 0; i < buffers.size(); ++i) {
       auto &buffer = buffers[i];
       construct(buffer.header.data());
@@ -936,10 +944,60 @@ void CheckAmprOrdering(Loader::SymbolDatabase &symbols, uint32_t file_id) {
       } else {
         construct_amm(buffer.header.data());
       }
-      Check(set_buffer(buffer.header.data(), buffer.data.data(), sizeof(buffer.data)) == OK,
+      Check(set_buffer(buffer.header.data(), buffer.data.data(), data_size) == OK,
             "initialize AMPR command buffer");
     }
   };
+  struct LiteralWrite { uint64_t value; uint32_t size; };
+  constexpr std::array literal_writes {
+      LiteralWrite{0, 8}, LiteralWrite{3, 8}, LiteralWrite{4, 12},
+      LiteralWrite{(uint64_t{1} << 34) - 1, 12},
+      LiteralWrite{uint64_t{1} << 34, 16}, LiteralWrite{UINT64_MAX, 16}};
+  for (const auto &literal : literal_writes) {
+    for (const bool immediately : std::array{false, true}) {
+      reset(1, literal.size);
+      auto *command_buffer = buffers[0].header.data();
+      uint64_t fence = ~literal.value;
+      const auto append = [&](uint64_t value) {
+        return immediately ? write_immediately(command_buffer, &fence, value, 1)
+                           : write_address(command_buffer, &fence, value);
+      };
+      Check(measure_write_address(&fence, literal.value) == literal.size &&
+                measure_write_immediately(&fence, literal.value, 1) == literal.size &&
+                append(literal.value) == OK && offset(command_buffer) == literal.size,
+            "literal write fits exactly its measured command size");
+      Check(append(~literal.value) == Libs::LibKernel::KERNEL_ERROR_EBUSY &&
+                offset(command_buffer) == literal.size,
+            "rejected write preserves the full command buffer offset");
+      Result result {1234, 5678};
+      uint32_t submission = 0;
+      Check(submit_apr(command_buffer, 3, &result, &submission) == OK &&
+                wait_apr(submission) == OK && result.result == OK &&
+                fence == literal.value,
+            "bounded write preserves all sixty-four literal bits");
+    }
+  }
+  reset(1, 768);
+  std::array<std::array<char, 3>, 16> batch_output {};
+  std::array<uint64_t, 16> batch_fences {};
+  auto *batch = buffers[0].header.data();
+  for (size_t i = 0; i < batch_fences.size(); ++i) {
+    Check(read_file(batch, reinterpret_cast<uint64_t>(&buffers[0].header[3]),
+                    reinterpret_cast<uint64_t>(&buffers[0].header[4]), file_id,
+                    batch_output[i].data(), batch_output[i].size(), 0) == OK &&
+              write_address(batch, &batch_fences[i], i + 1) == OK,
+          "sixteen APR read/fence pairs fit the guest's 768-byte command buffer");
+  }
+  Result batch_result {1234, 5678};
+  uint32_t batch_submission = 0;
+  Check(submit_apr(batch, 3, &batch_result, &batch_submission) == OK &&
+            wait_apr(batch_submission) == OK && batch_result.result == OK,
+        "complete the bounded APR read batch");
+  for (size_t i = 0; i < batch_fences.size(); ++i) {
+    Check(batch_fences[i] == i + 1 &&
+              std::memcmp(batch_output[i].data(), "APR", 3) == 0,
+          "every read publishes its data and completion, including requests fourteen and fifteen");
+  }
   // WaitCompare order: ==, unsigned >/<, !=, wrapped >=, signed >/<.
   struct Comparison { uint8_t compare; uint64_t blocked, reference, released; };
   constexpr std::array comparisons {