UPSTREAM: rmi4update: Optimize writing firmware blocks

Most HID devices allow appending the reflash command to the end of the firmware
block. This avoids sending a second report with the just the command. Also, after
the block is written HID devices send an attention report. Only read the F34 control
registers if waiting for that attention report times out.

Signed-off-by: Benson Leung <bleung@google.com>

Bug: 24809436

(cherry picked from commit 76743e425429076626df483691ce7abe563abd81)

Change-Id: I0a3276d77605843cbb6ddec221320a6048d7b925

2 files changed, 47 insertions, 29 deletions

diff --git a/rmi4update/rmi4update.cpp b/rmi4update/rmi4update.cpp
index 5faff1b..b846230 100644
--- a/rmi4update/rmi4update.cpp
+++ b/rmi4update/rmi4update.cpp
@@ -23,6 +23,7 @@
 #include <string.h>
 #include <stdlib.h>
 #include <errno.h>
+#include <alloca.h>
 
 #include "rmi4update.h"
 
@@ -425,6 +426,7 @@ int RMI4Update::WriteBlocks(unsigned char *block, unsigned short count, unsigned
 	unsigned char zeros[] = { 0, 0 };
 	int rc;
 	unsigned short addr;
+	unsigned char *blockWithCmd = (unsigned char *)alloca(m_blockSize + 1);
 
 	if (m_f34.GetFunctionVersion() == 0x1)
 		addr = m_f34.GetDataBase() + RMI_F34_BLOCK_DATA_V1_OFFSET;
@@ -436,19 +438,30 @@ int RMI4Update::WriteBlocks(unsigned char *block, unsigned short count, unsigned
 		return UPDATE_FAIL_WRITE_INITIAL_ZEROS;
 
 	for (blockNum = 0; blockNum < count; ++blockNum) {
-		rc = m_device.Write(addr, block, m_blockSize);
-		if (rc != m_blockSize) {
-			fprintf(stderr, "failed to write block %d\n", blockNum);
-			return UPDATE_FAIL_WRITE_BLOCK;
-		}
+		if (m_writeBlockWithCmd) {
+			memcpy(blockWithCmd, block, m_blockSize);
+			blockWithCmd[m_blockSize] = cmd;
+
+			rc = m_device.Write(addr, blockWithCmd, m_blockSize + 1);
+			if (rc != m_blockSize + 1) {
+				fprintf(stderr, "failed to write block %d\n", blockNum);
+				return UPDATE_FAIL_WRITE_BLOCK;
+			}
+		} else {
+			rc = m_device.Write(addr, block, m_blockSize);
+			if (rc != m_blockSize) {
+				fprintf(stderr, "failed to write block %d\n", blockNum);
+				return UPDATE_FAIL_WRITE_BLOCK;
+			}
 
-		rc = m_device.Write(m_f34StatusAddr, &cmd, 1);
-		if (rc != 1) {
-			fprintf(stderr, "failed to write command for block %d\n", blockNum);
-			return UPDATE_FAIL_WRITE_FLASH_COMMAND;
+			rc = m_device.Write(m_f34StatusAddr, &cmd, 1);
+			if (rc != 1) {
+				fprintf(stderr, "failed to write command for block %d\n", blockNum);
+				return UPDATE_FAIL_WRITE_FLASH_COMMAND;
+			}
 		}
 
-		rc = WaitForIdle(RMI_F34_IDLE_WAIT_MS);
+		rc = WaitForIdle(RMI_F34_IDLE_WAIT_MS, !m_writeBlockWithCmd);
 		if (rc != UPDATE_SUCCESS) {
 			fprintf(stderr, "failed to go into idle after writing block %d\n", blockNum);
 			return UPDATE_FAIL_NOT_IN_IDLE_STATE;
@@ -465,7 +478,7 @@ int RMI4Update::WriteBlocks(unsigned char *block, unsigned short count, unsigned
  * this will be true for HID, but other protocols will need to revert polling. Polling
  * is not implemented yet.
  */
-int RMI4Update::WaitForIdle(int timeout_ms)
+int RMI4Update::WaitForIdle(int timeout_ms, bool readF34OnSucess)
 {
 	int rc;
 	struct timeval tv;
@@ -486,24 +499,28 @@ int RMI4Update::WaitForIdle(int timeout_ms)
 			fprintf(stderr, "Timed out waiting for attn report\n");
 	}
 
-	rc = ReadF34Controls();
-	if (rc != UPDATE_SUCCESS)
-		return rc;
+	if (rc <= 0 || readF34OnSucess) {
+		rc = ReadF34Controls();
+		if (rc != UPDATE_SUCCESS)
+			return rc;
 
-	if (!m_f34Status && !m_f34Command) {
-		if (!m_programEnabled) {
-			fprintf(stderr, "Bootloader is idle but program_enabled bit isn't set.\n");
-			return UPDATE_FAIL_PROGRAMMING_NOT_ENABLED;
-		} else {
-			return UPDATE_SUCCESS;
+		if (!m_f34Status && !m_f34Command) {
+			if (!m_programEnabled) {
+				fprintf(stderr, "Bootloader is idle but program_enabled bit isn't set.\n");
+				return UPDATE_FAIL_PROGRAMMING_NOT_ENABLED;
+			} else {
+				return UPDATE_SUCCESS;
+			}
 		}
-	}
 
-	fprintf(stderr, "ERROR: Waiting for idle status.\n");
-	fprintf(stderr, "Command: %#04x\n", m_f34Command);
-	fprintf(stderr, "Status:  %#04x\n", m_f34Status);
-	fprintf(stderr, "Enabled: %d\n", m_programEnabled);
-	fprintf(stderr, "Idle:    %d\n", !m_f34Command && !m_f34Status);
+		fprintf(stderr, "ERROR: Waiting for idle status.\n");
+		fprintf(stderr, "Command: %#04x\n", m_f34Command);
+		fprintf(stderr, "Status:  %#04x\n", m_f34Status);
+		fprintf(stderr, "Enabled: %d\n", m_programEnabled);
+		fprintf(stderr, "Idle:    %d\n", !m_f34Command && !m_f34Status);
 
-	return UPDATE_FAIL_NOT_IN_IDLE_STATE;
+		return UPDATE_FAIL_NOT_IN_IDLE_STATE;
+	}
+
+	return UPDATE_SUCCESS;
 }
diff --git a/rmi4update/rmi4update.h b/rmi4update/rmi4update.h
index 7dff0f5..eaad0f6 100644
--- a/rmi4update/rmi4update.h
+++ b/rmi4update/rmi4update.h
@@ -27,7 +27,7 @@ class RMI4Update
 {
 public:
 	RMI4Update(RMIDevice & device, FirmwareImage & firmwareImage) : m_device(device), 
-			m_firmwareImage(firmwareImage)
+			m_firmwareImage(firmwareImage), m_writeBlockWithCmd(true)
 	{}
 	int UpdateFirmware(bool force = false, bool performLockdown = false);
 
@@ -39,7 +39,7 @@ private:
 	int WriteBootloaderID();
 	int EnterFlashProgramming();
 	int WriteBlocks(unsigned char *block, unsigned short count, unsigned char cmd);
-	int WaitForIdle(int timeout_ms);
+	int WaitForIdle(int timeout_ms, bool readF34OnSucess = true);
 	int GetFirmwareSize() { return m_blockSize * m_fwBlockCount; }
 	int GetConfigSize() { return m_blockSize * m_configBlockCount; }
 
@@ -52,6 +52,7 @@ private:
 
 	unsigned char m_deviceStatus;
 	unsigned char m_bootloaderID[RMI_BOOTLOADER_ID_SIZE];
+	bool m_writeBlockWithCmd;
 
 	/* F34 Controls */
 	unsigned char m_f34Command;