From b6378b2ee15acba1e4a37b1b836dded5e06fd4c6 Mon Sep 17 00:00:00 2001
From: Jay <jayhou@google.com>
Date: Wed, 17 Jan 2024 12:24:47 +0000
Subject: [PATCH] Call increment when new witness identity for both MMC and
 RPMB cases.

---
 trusted_os/flash.go |  6 +++---
 trusted_os/main.go  | 12 ++++++++++++
 trusted_os/rpmb.go  | 31 +++++++++++++++++++++++++++++++
 3 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/trusted_os/flash.go b/trusted_os/flash.go
index ae426bc..95947cd 100644
--- a/trusted_os/flash.go
+++ b/trusted_os/flash.go
@@ -181,8 +181,8 @@ func incrementWitnessIdentityMMC(card Card) error {
 	}
 
 	rBuf := bytes.NewReader(b)
-	var counter uint64
-	if err := binary.Read(rBuf, binary.LittleEndian, &counter); err != nil {
+	var counter uint32
+	if err := binary.Read(rBuf, binary.BigEndian, &counter); err != nil {
 		return err
 	}
 
@@ -191,7 +191,7 @@ func incrementWitnessIdentityMMC(card Card) error {
 
 	// Write
 	wBuf := new(bytes.Buffer)
-	if err := binary.Write(wBuf, binary.LittleEndian, counter); err != nil {
+	if err := binary.Write(wBuf, binary.BigEndian, counter); err != nil {
 		return err
 	}
 
diff --git a/trusted_os/main.go b/trusted_os/main.go
index f187ece..8056228 100644
--- a/trusted_os/main.go
+++ b/trusted_os/main.go
@@ -168,6 +168,18 @@ func main() {
 		log.Printf("Failed to determine OS MMC block (no OS installed?): %v", err)
 	}
 
+	newIdentity, err := newWitnessIdentity(Storage)
+	if err != nil {
+		log.Printf("Failed to read new witness identity MMC block: %v", err)
+	}
+	if newIdentity {
+		if false && imx6ul.SNVS.Available() {
+			rpmb.incrementWitnessIdentity()
+		} else {
+			incrementWitnessIdentityMMC(Storage)
+		}
+	}
+
 	log.Printf("SM log verification pub: %s", LogVerifier)
 	logVerifier, err := note.NewVerifier(LogVerifier)
 	if err != nil {
diff --git a/trusted_os/rpmb.go b/trusted_os/rpmb.go
index 007f9aa..431fe3c 100644
--- a/trusted_os/rpmb.go
+++ b/trusted_os/rpmb.go
@@ -37,12 +37,14 @@ import (
 const (
 	// RPMB sector for CVE-2020-13799 mitigation
 	dummySector = 0
+
 	// version epoch length
 	versionLength = 4
 	// RPMB sector for OS rollback protection
 	osVersionSector = 1
 	// RPMB sector for TA rollback protection
 	taVersionSector = 2
+
 	// RPMB sector for TA use
 	taUserSector = 3
 	// RPMB OTP flag bank
@@ -50,6 +52,11 @@ const (
 	// RPMB OTP flag word
 	rpmbFuseWord = 6
 
+	// witness identity counter length - uint32
+	witnessIdentityCounterLength = 4
+	// RPMB witness identity counter
+	rpmbWitnessIdentityCounter = 7
+
 	diversifierMAC = "ArmoryWitnessMAC"
 	iter           = 4096
 )
@@ -184,6 +191,30 @@ func (r *RPMB) checkVersion(offset uint16, s string) (err error) {
 	return
 }
 
+// incrementWitnessIdentity increments the counter in the RPMB area to
+// differentiate a new witness identity.
+func (r *RPMB) incrementWitnessIdentity() (err error) {
+	if r.partition == nil {
+		return errors.New("RPMB has not been initialized")
+	}
+
+	// Read
+	rBuf := make([]byte, witnessIdentityCounterLength)
+	if err = r.partition.Read(rpmbWitnessIdentityCounter, rBuf); err != nil {
+		return err
+	}
+	counter := binary.BigEndian.Uint32(rBuf)
+
+	// Increment
+	counter++
+
+	// Write
+	wBuf := make([]byte, witnessIdentityCounterLength)
+	binary.BigEndian.PutUint32(wBuf, counter)
+
+	return r.partition.Write(rpmbWitnessIdentityCounter, wBuf)
+}
+
 // transfer performs an authenticated data transfer to the card RPMB partition,
 // the input buffer can contain up to 256 bytes of data, n can be passed to
 // retrieve the partition write counter.