From e1fcae1719e7fa3082ad951a2044906d0db50ebc Mon Sep 17 00:00:00 2001
From: Jay <jayhou@google.com>
Date: Sun, 21 Jan 2024 12:21:51 +0000
Subject: [PATCH 01/16] Remove storage from RPMB struct.

---
 trusted_os/main.go          | 11 ++++++-----
 trusted_os/rpmb.go          | 28 ++++++++++++++++++----------
 trusted_os/storage.go       |  5 ++---
 trusted_os/storage_debug.go |  7 +++----
 4 files changed, 29 insertions(+), 22 deletions(-)

diff --git a/trusted_os/main.go b/trusted_os/main.go
index f187ece..357b668 100644
--- a/trusted_os/main.go
+++ b/trusted_os/main.go
@@ -134,13 +134,18 @@ func main() {
 	usbarmory.LED("blue", false)
 	usbarmory.LED("white", false)
 
-	Storage, rpmb := storage()
+	Storage := storage()
 	if Storage != nil {
 		if err = Storage.Detect(); err != nil {
 			log.Fatalf("SM failed to detect storage, %v", err)
 		}
 	}
 
+	rpmb, err := newRPMB(Storage)
+	if err != nil {
+		log.Fatalf("SM could not initialize rollback protection, %v", err)
+	}
+
 	rpc := &RPC{
 		RPMB:        rpmb,
 		Storage:     Storage,
@@ -155,10 +160,6 @@ func main() {
 	if false && imx6ul.SNVS.Available() {
 		log.Printf("SM version verification (%s)", Version)
 
-		if err = rpmb.init(); err != nil {
-			log.Fatalf("SM could not initialize rollback protection, %v", err)
-		}
-
 		if err = rpmb.checkVersion(osVersionSector, Version); err != nil {
 			log.Fatalf("SM firmware rollback check failure, %v", err)
 		}
diff --git a/trusted_os/rpmb.go b/trusted_os/rpmb.go
index 007f9aa..8b52ac6 100644
--- a/trusted_os/rpmb.go
+++ b/trusted_os/rpmb.go
@@ -55,11 +55,10 @@ const (
 )
 
 type RPMB struct {
-	storage   *usdhc.USDHC
 	partition *rpmb.RPMB
 }
 
-func (r *RPMB) init() (err error) {
+func newRPMB(storage Card) (r *RPMB, err error) {
 	// derived key for RPBM MAC generation
 	var dk []byte
 
@@ -73,23 +72,32 @@ func (r *RPMB) init() (err error) {
 	}
 
 	if err != nil {
-		return fmt.Errorf("could not derive RPMB key (%v)", err)
+		return nil, fmt.Errorf("could not derive RPMB key (%v)", err)
 	}
 
 	uid := imx6ul.UniqueID()
 
-	// setup RPMB partition
+	card, ok := storage.(*usdhc.USDHC)
+	if !ok {
+		return nil, errors.New("could not assert type *usdhc.USDHC from Card")
+	}
+
+	// setup RPMB
+	r = &RPMB{}
 	r.partition, err = rpmb.Init(
-		r.storage,
+		card,
 		pbkdf2.Key(dk, uid[:], iter, sha256.Size, sha256.New),
 		dummySector,
 	)
+	if err != nil {
+		return nil, errors.New("could not init RPMB TODO")
+	}
 
 	var e *rpmb.OperationError
 	_, err = r.partition.Counter(false)
 
 	if !(errors.As(err, &e) && e.Result == rpmb.AuthenticationKeyNotYetProgrammed) {
-		return
+		return nil, fmt.Errorf("RPMB could not be initialized: %v", err)
 	}
 
 	// Fuse a bit to indicate previous key programming to prevent malicious
@@ -97,20 +105,20 @@ func (r *RPMB) init() (err error) {
 	//
 	// If already fused refuse to do any programming and bail.
 	if res, err := otp.ReadOCOTP(rpmbFuseBank, rpmbFuseWord, 0, 1); err != nil || bytes.Equal(res, []byte{1}) {
-		return fmt.Errorf("could not read RPMB program key flag (%x, %v)", res, err)
+		return nil, fmt.Errorf("could not read RPMB program key flag (%x, %v)", res, err)
 	}
 
 	if err = otp.BlowOCOTP(rpmbFuseBank, rpmbFuseWord, 0, 1, []byte{1}); err != nil {
-		return fmt.Errorf("could not fuse RPMB program key flag (%v)", err)
+		return nil, fmt.Errorf("could not fuse RPMB program key flag (%v)", err)
 	}
 
 	log.Print("RPMB authentication key not yet programmed, programming")
 
 	if err = r.partition.ProgramKey(); err != nil {
-		return fmt.Errorf("could not program RPMB key")
+		return nil, fmt.Errorf("could not program RPMB key")
 	}
 
-	return
+	return r, nil
 }
 
 func parseVersion(s string) (version uint32, err error) {
diff --git a/trusted_os/storage.go b/trusted_os/storage.go
index 4c16a0a..4b1a17c 100644
--- a/trusted_os/storage.go
+++ b/trusted_os/storage.go
@@ -21,7 +21,6 @@ import (
 	usbarmory "github.com/usbarmory/tamago/board/usbarmory/mk2"
 )
 
-func storage() (Card, *RPMB) {
-	s := usbarmory.MMC
-	return s, &RPMB{storage: s}
+func storage() Card {
+	return usbarmory.MMC
 }
diff --git a/trusted_os/storage_debug.go b/trusted_os/storage_debug.go
index f5cda2b..6345196 100644
--- a/trusted_os/storage_debug.go
+++ b/trusted_os/storage_debug.go
@@ -35,12 +35,11 @@ const (
 
 // storage will return MMC backed storage if running on real hardware, or
 // a fake in-memory storage device otherwise.
-func storage() (Card, *RPMB) {
+func storage() Card {
 	if imx6ul.Native {
-		s := usbarmory.MMC
-		return s, &RPMB{storage: s}
+		return usbarmory.MMC
 	}
-	return newFakeCard(fakeCardNumBlocks), &RPMB{}
+	return newFakeCard(fakeCardNumBlocks)
 }
 
 // fakeCard is an implementation of an in-memory storage device.

From deff3c0afc32b00ed07410a7057e038810474577 Mon Sep 17 00:00:00 2001
From: Jay <jayhou@google.com>
Date: Fri, 26 Jan 2024 13:18:57 +0000
Subject: [PATCH 02/16] Add fake RPMB.

---
 Makefile                |   3 +
 rpmb/op.go              |  10 ++--
 trusted_os/rpmb.go      |   5 +-
 trusted_os/rpmb_fake.go | 130 ++++++++++++++++++++++++++++++++++++++++
 4 files changed, 142 insertions(+), 6 deletions(-)
 create mode 100644 trusted_os/rpmb_fake.go

diff --git a/Makefile b/Makefile
index 4123ea7..c19fdc2 100644
--- a/Makefile
+++ b/Makefile
@@ -30,6 +30,9 @@ endif
 ifeq ("${FAKE_STORAGE}", "1")
 	BUILD_TAGS := ${BUILD_TAGS},fake_storage
 endif
+ifeq ("${FAKE_RPMB}", "1")
+	BUILD_TAGS := ${BUILD_TAGS},fake_rpmb
+endif
 
 APP := ""
 TEXT_START = 0x80010000 # ramStart (defined in mem.go under relevant tamago/soc package) + 0x10000
diff --git a/rpmb/op.go b/rpmb/op.go
index f7f7ab3..bbf4b10 100644
--- a/rpmb/op.go
+++ b/rpmb/op.go
@@ -26,7 +26,7 @@ import (
 )
 
 const (
-	frameLength = 512
+	FrameLength = 512
 	macOffset   = 284
 )
 
@@ -112,7 +112,7 @@ func (p *RPMB) op(req *DataFrame, cfg *Config) (res *DataFrame, err error) {
 	mac := hmac.New(sha256.New, p.key[:])
 
 	if cfg.RequestMAC {
-		mac.Write(req.Bytes()[frameLength-macOffset:])
+		mac.Write(req.Bytes()[FrameLength-macOffset:])
 		copy(req.KeyMAC[:], mac.Sum(nil))
 		mac.Reset()
 	}
@@ -145,7 +145,7 @@ func (p *RPMB) op(req *DataFrame, cfg *Config) (res *DataFrame, err error) {
 		}
 	}
 
-	buf := make([]byte, frameLength)
+	buf := make([]byte, FrameLength)
 	res = &DataFrame{}
 
 	// read response
@@ -161,7 +161,7 @@ func (p *RPMB) op(req *DataFrame, cfg *Config) (res *DataFrame, err error) {
 	// validate response
 
 	if cfg.ResponseMAC {
-		mac.Write(buf[frameLength-macOffset:])
+		mac.Write(buf[FrameLength-macOffset:])
 
 		if !hmac.Equal(res.KeyMAC[:], mac.Sum(nil)) {
 			return nil, errors.New("invalid response MAC")
@@ -186,7 +186,7 @@ func (p *RPMB) op(req *DataFrame, cfg *Config) (res *DataFrame, err error) {
 }
 
 func (p *RPMB) transfer(kind byte, offset uint16, buf []byte) (err error) {
-	if len(buf) > frameLength/2 {
+	if len(buf) > FrameLength/2 {
 		return errors.New("transfer size must not exceed 256 bytes")
 	}
 
diff --git a/trusted_os/rpmb.go b/trusted_os/rpmb.go
index 8b52ac6..c7b1cea 100644
--- a/trusted_os/rpmb.go
+++ b/trusted_os/rpmb.go
@@ -12,6 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+//go:build !fake_rpmb
+// +build !fake_rpmb
+
 package main
 
 import (
@@ -90,7 +93,7 @@ func newRPMB(storage Card) (r *RPMB, err error) {
 		dummySector,
 	)
 	if err != nil {
-		return nil, errors.New("could not init RPMB TODO")
+		return nil, fmt.Errorf("RPMB could not be initialized: %v", err)
 	}
 
 	var e *rpmb.OperationError
diff --git a/trusted_os/rpmb_fake.go b/trusted_os/rpmb_fake.go
new file mode 100644
index 0000000..06801c2
--- /dev/null
+++ b/trusted_os/rpmb_fake.go
@@ -0,0 +1,130 @@
+// Copyright 2022 The Armored Witness OS authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build fake_rpmb
+// +build fake_rpmb
+
+package main
+
+import (
+	"encoding/binary"
+	"errors"
+	"strconv"
+
+	"github.com/transparency-dev/armored-witness-os/rpmb"
+)
+
+const (
+	// 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
+
+	sectorLength = 256
+	memoryLength = 1024
+)
+
+type RPMB struct {
+	mem     [memoryLength]byte
+	counter uint32
+}
+
+func newRPMB(_ Card) (*RPMB, error) {
+	return &RPMB{}, nil
+}
+
+func parseVersion(s string) (version uint32, err error) {
+	v, err := strconv.Atoi(s)
+
+	if err != nil {
+		return
+	}
+
+	return uint32(v), nil
+}
+
+// expectedVersion returns the version epoch stored in a fake RPMB area.
+func (r *RPMB) expectedVersion(sector uint16) (version uint32, err error) {
+	buf := make([]byte, versionLength)
+	copy(buf, r.mem[sector*sectorLength:])
+
+	return binary.BigEndian.Uint32(buf), nil
+}
+
+// updateVersion writes a new version epoch in a fake RPMB area.
+func (r *RPMB) updateVersion(sector uint16, version uint32) (err error) {
+	buf := make([]byte, versionLength)
+	binary.BigEndian.PutUint32(buf, version)
+
+	copy(r.mem[sector*sectorLength:], buf)
+	r.counter++
+
+	return nil
+}
+
+// checkVersion verifies version information against fake RPMB stored data.
+//
+// If the passed version is older than the RPMB area information of the
+// internal eMMC an error is returned.
+//
+// If the passed version is more recent than the RPMB area information then the
+// internal eMMC is updated with it.
+func (r *RPMB) checkVersion(sector uint16, s string) (err error) {
+	version, err := parseVersion(s)
+
+	if err != nil {
+		return
+	}
+
+	expectedVersion, err := r.expectedVersion(sector)
+
+	if err != nil {
+		return
+	}
+
+	switch {
+	case expectedVersion > version:
+		return errors.New("version mismatch")
+	case expectedVersion == version:
+		return
+	case expectedVersion < version:
+		return r.updateVersion(sector, version)
+	}
+
+	return
+}
+
+// transfer performs a data transfer to the fake RPMB area,
+// the input buffer can contain up to 256 bytes of data, n can be passed to
+// retrieve the write counter.
+func (r *RPMB) transfer(sector uint16, buf []byte, n *uint32, write bool) (err error) {
+	if len(buf) > rpmb.FrameLength/2 {
+		return errors.New("transfer size must not exceed 256 bytes")
+	}
+
+	if write {
+		copy(r.mem[sector*sectorLength:], buf)
+		r.counter++
+	} else {
+		copy(buf, r.mem[sector*sectorLength:])
+	}
+
+	*n = r.counter
+	return
+}

From d90d8e5787d435843294d52fdc0fddc6063abb75 Mon Sep 17 00:00:00 2001
From: Jay <jayhou@google.com>
Date: Mon, 29 Jan 2024 12:44:01 +0000
Subject: [PATCH 03/16] Make counter read optional for RPMB.

---
 trusted_os/rpmb_fake.go | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/trusted_os/rpmb_fake.go b/trusted_os/rpmb_fake.go
index 06801c2..72f6001 100644
--- a/trusted_os/rpmb_fake.go
+++ b/trusted_os/rpmb_fake.go
@@ -125,6 +125,8 @@ func (r *RPMB) transfer(sector uint16, buf []byte, n *uint32, write bool) (err e
 		copy(buf, r.mem[sector*sectorLength:])
 	}
 
-	*n = r.counter
+	if n != nil {
+		*n = r.counter
+	}
 	return
 }

From 1a0d19e08eb352197afdcf69404b6ecf08e2d0e7 Mon Sep 17 00:00:00 2001
From: Jay <jayhou@google.com>
Date: Mon, 29 Jan 2024 22:52:12 +0000
Subject: [PATCH 04/16] Use 2D memory array for fake RPMB.

---
 trusted_os/rpmb_fake.go | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/trusted_os/rpmb_fake.go b/trusted_os/rpmb_fake.go
index 72f6001..f8cd8c2 100644
--- a/trusted_os/rpmb_fake.go
+++ b/trusted_os/rpmb_fake.go
@@ -37,11 +37,11 @@ const (
 	taUserSector = 3
 
 	sectorLength = 256
-	memoryLength = 1024
+	numSectors   = 16
 )
 
 type RPMB struct {
-	mem     [memoryLength]byte
+	mem     [numSectors][sectorLength]byte
 	counter uint32
 }
 
@@ -62,7 +62,7 @@ func parseVersion(s string) (version uint32, err error) {
 // expectedVersion returns the version epoch stored in a fake RPMB area.
 func (r *RPMB) expectedVersion(sector uint16) (version uint32, err error) {
 	buf := make([]byte, versionLength)
-	copy(buf, r.mem[sector*sectorLength:])
+	copy(buf, r.mem[sector])
 
 	return binary.BigEndian.Uint32(buf), nil
 }
@@ -72,7 +72,7 @@ func (r *RPMB) updateVersion(sector uint16, version uint32) (err error) {
 	buf := make([]byte, versionLength)
 	binary.BigEndian.PutUint32(buf, version)
 
-	copy(r.mem[sector*sectorLength:], buf)
+	copy(r.mem[sector], buf)
 	r.counter++
 
 	return nil
@@ -119,10 +119,10 @@ func (r *RPMB) transfer(sector uint16, buf []byte, n *uint32, write bool) (err e
 	}
 
 	if write {
-		copy(r.mem[sector*sectorLength:], buf)
+		copy(r.mem[sector], buf)
 		r.counter++
 	} else {
-		copy(buf, r.mem[sector*sectorLength:])
+		copy(buf, r.mem[sector])
 	}
 
 	if n != nil {

From cebad448c3b18dbaa0abd58cdd931b708d9b2dd5 Mon Sep 17 00:00:00 2001
From: Jay <jayhou@google.com>
Date: Tue, 16 Jan 2024 13:26:30 +0000
Subject: [PATCH 05/16] Add sector constants and read/write for MMC blocks
 related to new witness identity.

---
 trusted_os/flash.go | 84 ++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 76 insertions(+), 8 deletions(-)

diff --git a/trusted_os/flash.go b/trusted_os/flash.go
index 738a60d..ae426bc 100644
--- a/trusted_os/flash.go
+++ b/trusted_os/flash.go
@@ -15,6 +15,8 @@
 package main
 
 import (
+	"bytes"
+	"encoding/binary"
 	"errors"
 	"fmt"
 	"log"
@@ -34,14 +36,25 @@ import (
 
 const (
 	expectedBlockSize = 512 // Expected size of MMC block in bytes
-	otaLimit          = 31457280
-	taConfBlock       = 0x200000
-	taBlockA          = 0x200050
-	taBlockB          = 0x2FD050
-	osConfBlock       = 0x5000
-	osBlockA          = 0x5050
-	osBlockB          = 0x102828
-	batchSize         = 2048
+
+	otaLimit = 31457280
+
+	taConfBlock = 0x200000
+	taBlockA    = 0x200050
+	taBlockB    = 0x2FD050
+
+	osConfBlock = 0x5000
+	osBlockA    = 0x5050
+	osBlockB    = 0x102828
+	// This sector is interpreted as whether the device should wipe its previous
+	// data and boot as a new witness identity.
+	newIdentityBlock = 0x1FFF00
+	// This sector contains a counter to differentiate between witness
+	// identities. For production devices, this counter should live in RPMB,
+	// so that it is secure, read-only, and overflow protected.
+	mmcIdentityCounter = 0x1FFF01
+
+	batchSize = 2048
 )
 
 const (
@@ -133,6 +146,61 @@ func determineLoadedOSBlock(card Card) error {
 	return nil
 }
 
+// newWitnessIdentity reads the newIdentityBlock from MMC.
+func newWitnessIdentity(card Card) (bool, error) {
+	blockSize := card.Info().BlockSize
+	if blockSize != expectedBlockSize {
+		return false, fmt.Errorf("h/w invariant error - expected MMC blocksize %d, found %d", expectedBlockSize, blockSize)
+	}
+
+	b, err := card.Read(newIdentityBlock*expectedBlockSize, 1)
+	if err != nil {
+		return false, err
+	}
+
+	buf := bytes.NewReader(b)
+	var newIdentity bool
+	if err := binary.Read(buf, binary.LittleEndian, &newIdentity); err != nil {
+		return false, err
+	}
+
+	return newIdentity, nil
+}
+
+// incrementWitnessIdentityMMC increments the mmcIdentityCounter in the MMC block.
+func incrementWitnessIdentityMMC(card Card) error {
+	blockSize := card.Info().BlockSize
+	if blockSize != expectedBlockSize {
+		return fmt.Errorf("h/w invariant error - expected MMC blocksize %d, found %d", expectedBlockSize, blockSize)
+	}
+
+	// Read
+	b, err := card.Read(mmcIdentityCounter*expectedBlockSize, 1)
+	if err != nil {
+		return err
+	}
+
+	rBuf := bytes.NewReader(b)
+	var counter uint64
+	if err := binary.Read(rBuf, binary.LittleEndian, &counter); err != nil {
+		return err
+	}
+
+	// Increment
+	counter++
+
+	// Write
+	wBuf := new(bytes.Buffer)
+	if err := binary.Write(wBuf, binary.LittleEndian, counter); err != nil {
+		return err
+	}
+
+	if err := card.WriteBlocks(mmcIdentityCounter*expectedBlockSize, wBuf.Bytes()); err != nil {
+		return err
+	}
+	return nil
+}
+
 // read reads the trusted applet bundle from internal storage, the
 // applet and FT proofs are *not* verified by this function.
 //

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

---
 trusted_os/flash.go |  8 ++++----
 trusted_os/main.go  | 13 +++++++++++++
 trusted_os/rpmb.go  | 31 +++++++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/trusted_os/flash.go b/trusted_os/flash.go
index ae426bc..e5108ff 100644
--- a/trusted_os/flash.go
+++ b/trusted_os/flash.go
@@ -160,7 +160,7 @@ func newWitnessIdentity(card Card) (bool, error) {
 
 	buf := bytes.NewReader(b)
 	var newIdentity bool
-	if err := binary.Read(buf, binary.LittleEndian, &newIdentity); err != nil {
+	if err := binary.Read(buf, binary.BigEndian, &newIdentity); err != nil {
 		return false, err
 	}
 
@@ -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 357b668..f266f7b 100644
--- a/trusted_os/main.go
+++ b/trusted_os/main.go
@@ -169,6 +169,19 @@ 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 {
+		// TODO: disable for now
+		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 c7b1cea..2a90240 100644
--- a/trusted_os/rpmb.go
+++ b/trusted_os/rpmb.go
@@ -40,12 +40,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
@@ -53,6 +55,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
 )
@@ -195,6 +202,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.

From e4da79cda9c26e3b985ac0533fe5fd9babdb3844 Mon Sep 17 00:00:00 2001
From: Jay <jayhou@google.com>
Date: Thu, 18 Jan 2024 11:43:41 +0000
Subject: [PATCH 07/16] Add RPCs for reading the witness identity counter.

---
 trusted_os/rpc.go | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/trusted_os/rpc.go b/trusted_os/rpc.go
index da87577..1908878 100644
--- a/trusted_os/rpc.go
+++ b/trusted_os/rpc.go
@@ -15,8 +15,10 @@
 package main
 
 import (
+	"bytes"
 	"crypto/aes"
 	"crypto/sha256"
+	"encoding/binary"
 	"errors"
 	"log"
 	"net"
@@ -156,6 +158,26 @@ func (r *RPC) Read(xfer rpc.Read, out *[]byte) (err error) {
 	return
 }
 
+// ReadIdentityCounterStorage reads data from the storage media sector
+// containing the witness identity counter.
+func (r *RPC) ReadIdentityCounterStorage(_ any, counter *uint32) error {
+	if r.Storage == nil {
+		return errors.New("missing Storage")
+	}
+
+	b, err := r.Storage.Read(mmcIdentityCounter, 1)
+	if err != nil {
+		return err
+	}
+
+	buf := bytes.NewReader(b)
+	if err := binary.Read(buf, binary.BigEndian, *counter); err != nil {
+		return err
+	}
+
+	return err
+}
+
 // WriteRPMB performs an authenticated data transfer to the card RPMB partition
 // sector allocated to the Trusted Applet. The input buffer can contain up to
 // 256 bytes of data, n can be passed to retrieve the partition write counter.
@@ -171,6 +193,19 @@ func (r *RPC) ReadRPMB(buf []byte, n *uint32) error {
 	return r.RPMB.transfer(taUserSector, buf, n, false)
 }
 
+// ReadIdentityCounterRPMB performs an authenticated data transfer from the card RPMB
+// partition sector allocated to the witness identity counter. It returns the
+// value stored in this area.
+func (r *RPC) ReadIdentityCounterRPMB(_ any, counter *uint32) error {
+	buf := make([]byte, witnessIdentityCounterLength)
+	if err := r.RPMB.transfer(rpmbWitnessIdentityCounter, buf, nil, false); err != nil {
+		return err
+	}
+
+	*counter = binary.BigEndian.Uint32(buf)
+	return nil
+}
+
 // DeriveKey derives a hardware unique key in a manner equivalent to PKCS#11
 // C_DeriveKey with CKM_AES_CBC_ENCRYPT_DATA.
 //

From 5a2c38c98d4dd66c6f02bec5033b44d12898a7d2 Mon Sep 17 00:00:00 2001
From: Jay <jayhou@google.com>
Date: Mon, 29 Jan 2024 23:14:07 +0000
Subject: [PATCH 08/16] Use fake RPMB to implement the counter.

---
 trusted_os/flash.go     | 64 -----------------------------------------
 trusted_os/main.go      | 13 ---------
 trusted_os/rpc.go       | 23 +--------------
 trusted_os/rpmb.go      | 41 +++++++++++++-------------
 trusted_os/rpmb_fake.go | 23 +++++++++++++++
 5 files changed, 44 insertions(+), 120 deletions(-)

diff --git a/trusted_os/flash.go b/trusted_os/flash.go
index e5108ff..d303930 100644
--- a/trusted_os/flash.go
+++ b/trusted_os/flash.go
@@ -15,8 +15,6 @@
 package main
 
 import (
-	"bytes"
-	"encoding/binary"
 	"errors"
 	"fmt"
 	"log"
@@ -46,13 +44,6 @@ const (
 	osConfBlock = 0x5000
 	osBlockA    = 0x5050
 	osBlockB    = 0x102828
-	// This sector is interpreted as whether the device should wipe its previous
-	// data and boot as a new witness identity.
-	newIdentityBlock = 0x1FFF00
-	// This sector contains a counter to differentiate between witness
-	// identities. For production devices, this counter should live in RPMB,
-	// so that it is secure, read-only, and overflow protected.
-	mmcIdentityCounter = 0x1FFF01
 
 	batchSize = 2048
 )
@@ -146,61 +137,6 @@ func determineLoadedOSBlock(card Card) error {
 	return nil
 }
 
-// newWitnessIdentity reads the newIdentityBlock from MMC.
-func newWitnessIdentity(card Card) (bool, error) {
-	blockSize := card.Info().BlockSize
-	if blockSize != expectedBlockSize {
-		return false, fmt.Errorf("h/w invariant error - expected MMC blocksize %d, found %d", expectedBlockSize, blockSize)
-	}
-
-	b, err := card.Read(newIdentityBlock*expectedBlockSize, 1)
-	if err != nil {
-		return false, err
-	}
-
-	buf := bytes.NewReader(b)
-	var newIdentity bool
-	if err := binary.Read(buf, binary.BigEndian, &newIdentity); err != nil {
-		return false, err
-	}
-
-	return newIdentity, nil
-}
-
-// incrementWitnessIdentityMMC increments the mmcIdentityCounter in the MMC block.
-func incrementWitnessIdentityMMC(card Card) error {
-	blockSize := card.Info().BlockSize
-	if blockSize != expectedBlockSize {
-		return fmt.Errorf("h/w invariant error - expected MMC blocksize %d, found %d", expectedBlockSize, blockSize)
-	}
-
-	// Read
-	b, err := card.Read(mmcIdentityCounter*expectedBlockSize, 1)
-	if err != nil {
-		return err
-	}
-
-	rBuf := bytes.NewReader(b)
-	var counter uint32
-	if err := binary.Read(rBuf, binary.BigEndian, &counter); err != nil {
-		return err
-	}
-
-	// Increment
-	counter++
-
-	// Write
-	wBuf := new(bytes.Buffer)
-	if err := binary.Write(wBuf, binary.BigEndian, counter); err != nil {
-		return err
-	}
-
-	if err := card.WriteBlocks(mmcIdentityCounter*expectedBlockSize, wBuf.Bytes()); err != nil {
-		return err
-	}
-	return nil
-}
-
 // read reads the trusted applet bundle from internal storage, the
 // applet and FT proofs are *not* verified by this function.
 //
diff --git a/trusted_os/main.go b/trusted_os/main.go
index f266f7b..357b668 100644
--- a/trusted_os/main.go
+++ b/trusted_os/main.go
@@ -169,19 +169,6 @@ 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 {
-		// TODO: disable for now
-		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/rpc.go b/trusted_os/rpc.go
index 1908878..616b630 100644
--- a/trusted_os/rpc.go
+++ b/trusted_os/rpc.go
@@ -15,7 +15,6 @@
 package main
 
 import (
-	"bytes"
 	"crypto/aes"
 	"crypto/sha256"
 	"encoding/binary"
@@ -158,26 +157,6 @@ func (r *RPC) Read(xfer rpc.Read, out *[]byte) (err error) {
 	return
 }
 
-// ReadIdentityCounterStorage reads data from the storage media sector
-// containing the witness identity counter.
-func (r *RPC) ReadIdentityCounterStorage(_ any, counter *uint32) error {
-	if r.Storage == nil {
-		return errors.New("missing Storage")
-	}
-
-	b, err := r.Storage.Read(mmcIdentityCounter, 1)
-	if err != nil {
-		return err
-	}
-
-	buf := bytes.NewReader(b)
-	if err := binary.Read(buf, binary.BigEndian, *counter); err != nil {
-		return err
-	}
-
-	return err
-}
-
 // WriteRPMB performs an authenticated data transfer to the card RPMB partition
 // sector allocated to the Trusted Applet. The input buffer can contain up to
 // 256 bytes of data, n can be passed to retrieve the partition write counter.
@@ -198,7 +177,7 @@ func (r *RPC) ReadRPMB(buf []byte, n *uint32) error {
 // value stored in this area.
 func (r *RPC) ReadIdentityCounterRPMB(_ any, counter *uint32) error {
 	buf := make([]byte, witnessIdentityCounterLength)
-	if err := r.RPMB.transfer(rpmbWitnessIdentityCounter, buf, nil, false); err != nil {
+	if err := r.RPMB.transfer(witnessIdentityCounter, buf, nil, false); err != nil {
 		return err
 	}
 
diff --git a/trusted_os/rpmb.go b/trusted_os/rpmb.go
index 2a90240..026205f 100644
--- a/trusted_os/rpmb.go
+++ b/trusted_os/rpmb.go
@@ -58,7 +58,7 @@ const (
 	// witness identity counter length - uint32
 	witnessIdentityCounterLength = 4
 	// RPMB witness identity counter
-	rpmbWitnessIdentityCounter = 7
+	witnessIdentityCounter = 7
 
 	diversifierMAC = "ArmoryWitnessMAC"
 	iter           = 4096
@@ -202,6 +202,23 @@ func (r *RPMB) checkVersion(offset uint16, s string) (err error) {
 	return
 }
 
+// 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.
+func (r *RPMB) transfer(offset uint16, buf []byte, n *uint32, write bool) (err error) {
+	if write {
+		err = r.partition.Write(offset, buf)
+	} else {
+		err = r.partition.Read(offset, buf)
+	}
+
+	if err != nil && n != nil {
+		*n, err = r.partition.Counter(true)
+	}
+
+	return
+}
+
 // incrementWitnessIdentity increments the counter in the RPMB area to
 // differentiate a new witness identity.
 func (r *RPMB) incrementWitnessIdentity() (err error) {
@@ -211,7 +228,7 @@ func (r *RPMB) incrementWitnessIdentity() (err error) {
 
 	// Read
 	rBuf := make([]byte, witnessIdentityCounterLength)
-	if err = r.partition.Read(rpmbWitnessIdentityCounter, rBuf); err != nil {
+	if err = r.partition.Read(witnessIdentityCounter, rBuf); err != nil {
 		return err
 	}
 	counter := binary.BigEndian.Uint32(rBuf)
@@ -222,23 +239,5 @@ func (r *RPMB) incrementWitnessIdentity() (err error) {
 	// 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.
-func (r *RPMB) transfer(offset uint16, buf []byte, n *uint32, write bool) (err error) {
-	if write {
-		err = r.partition.Write(offset, buf)
-	} else {
-		err = r.partition.Read(offset, buf)
-	}
-
-	if err != nil && n != nil {
-		*n, err = r.partition.Counter(true)
-	}
-
-	return
+	return r.partition.Write(witnessIdentityCounter, wBuf)
 }
diff --git a/trusted_os/rpmb_fake.go b/trusted_os/rpmb_fake.go
index f8cd8c2..4d6f22a 100644
--- a/trusted_os/rpmb_fake.go
+++ b/trusted_os/rpmb_fake.go
@@ -36,6 +36,10 @@ const (
 	// RPMB sector for TA use
 	taUserSector = 3
 
+	witnessIdentityCounterLength = 4
+	// RPMB sector for witness identity counter
+	witnessIdentityCounterSector = 4
+
 	sectorLength = 256
 	numSectors   = 16
 )
@@ -130,3 +134,22 @@ func (r *RPMB) transfer(sector uint16, buf []byte, n *uint32, write bool) (err e
 	}
 	return
 }
+
+// incrementWitnessIdentity increments the counter in the RPMB area to
+// differentiate a new witness identity.
+func (r *RPMB) incrementWitnessIdentity() error {
+	// Read
+	rBuf := make([]byte, witnessIdentityCounterLength)
+	if err := r.transfer(witnessIdentityCounterSector, rBuf, nil, false); err != nil {
+		return err
+	}
+	counter := binary.BigEndian.Uint32(rBuf)
+
+	// Increment
+	counter++
+
+	// Write
+	wBuf := make([]byte, witnessIdentityCounterLength)
+	binary.BigEndian.PutUint32(wBuf, counter)
+	return r.transfer(witnessIdentityCounterSector, wBuf, nil, true)
+}

From ad2ae61285a62493d78d9bc38cafb48ec2990929 Mon Sep 17 00:00:00 2001
From: Jay <jayhou@google.com>
Date: Tue, 30 Jan 2024 11:28:34 +0000
Subject: [PATCH 09/16] Add identity counter to witness status data structures.

---
 api/api.go     | 19 +++++-----
 api/api.pb.go  | 97 +++++++++++++++++++++++++++-----------------------
 api/api.proto  | 92 ++++++++++++++++++++++++-----------------------
 api/rpc/rpc.go |  3 ++
 4 files changed, 113 insertions(+), 98 deletions(-)

diff --git a/api/api.go b/api/api.go
index 8b7c241..fe71a31 100644
--- a/api/api.go
+++ b/api/api.go
@@ -92,17 +92,18 @@ func (p *Status) Print() string {
 	var status bytes.Buffer
 
 	status.WriteString("----------------------------------------------------------- Trusted OS ----\n")
-	status.WriteString(fmt.Sprintf("Serial number ..........: %s\n", p.Serial))
-	status.WriteString(fmt.Sprintf("Secure Boot ............: %v\n", p.HAB))
-	status.WriteString(fmt.Sprintf("Revision ...............: %s\n", p.Revision))
-	status.WriteString(fmt.Sprintf("Version ................: %d (%s)\n", p.Version, time.Unix(int64(p.Version), 0)))
-	status.WriteString(fmt.Sprintf("Runtime ................: %s\n", p.Runtime))
-	status.WriteString(fmt.Sprintf("Link ...................: %v\n", p.Link))
+	status.WriteString(fmt.Sprintf("Serial number ..............: %s\n", p.Serial))
+	status.WriteString(fmt.Sprintf("Secure Boot ................: %v\n", p.HAB))
+	status.WriteString(fmt.Sprintf("Revision ...................: %s\n", p.Revision))
+	status.WriteString(fmt.Sprintf("Version ....................: %d (%s)\n", p.Version, time.Unix(int64(p.Version), 0)))
+	status.WriteString(fmt.Sprintf("Runtime ....................: %s\n", p.Runtime))
+	status.WriteString(fmt.Sprintf("Link .......................: %v\n", p.Link))
 	if p.Witness != nil {
-		status.WriteString(fmt.Sprintf("Witness/Identity .......: %v\n", p.Witness.Identity))
-		status.WriteString(fmt.Sprintf("Witness/IP .............: %v", p.Witness.IP))
+		status.WriteString(fmt.Sprintf("Witness/Identity ...........: %v\n", p.Witness.Identity))
+		status.WriteString(fmt.Sprintf("Witness/IP .................: %v", p.Witness.IP))
+		status.WriteString(fmt.Sprintf("Witness/IdentityCounter ....: %d", p.Witness.IdentityCounter))
 	} else {
-		status.WriteString(fmt.Sprint("Witness ................: <no status>"))
+		status.WriteString(fmt.Sprint("Witness ....................: <no status>"))
 	}
 
 	return status.String()
diff --git a/api/api.pb.go b/api/api.pb.go
index c6af91b..0aed4b4 100644
--- a/api/api.pb.go
+++ b/api/api.pb.go
@@ -16,7 +16,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.28.1
-// 	protoc        v3.20.1
+// 	protoc        v3.21.12
 // source: api.proto
 
 package api
@@ -83,17 +83,14 @@ func (ErrorCode) EnumDescriptor() ([]byte, []int) {
 	return file_api_proto_rawDescGZIP(), []int{0}
 }
 
+// AppletUpdate
 //
+// A `AppletUpdate` represents an OTA sequence applet slice.
 //
-//AppletUpdate
-//
-//A `AppletUpdate` represents an OTA sequence applet slice.
-//
-//The `TotalChunks` value indicates the total number of chunks for the update,
-//`Seq` is the transmitted AppletUpdate chunk number:
-//- `0` indicates that the struct contains verification data in `Header`.
-//- `1` onwards identifies the first, second, ... chunk of firmware image data.
-//
+// The `TotalChunks` value indicates the total number of chunks for the update,
+// `Seq` is the transmitted AppletUpdate chunk number:
+// - `0` indicates that the struct contains verification data in `Header`.
+// - `1` onwards identifies the first, second, ... chunk of firmware image data.
 type AppletUpdate struct {
 	state         protoimpl.MessageState
 	sizeCache     protoimpl.SizeCache
@@ -102,6 +99,7 @@ type AppletUpdate struct {
 	Total uint32 `protobuf:"varint,1,opt,name=Total,proto3" json:"Total,omitempty"`
 	Seq   uint32 `protobuf:"varint,2,opt,name=Seq,proto3" json:"Seq,omitempty"`
 	// Types that are assignable to Payload:
+	//
 	//	*AppletUpdate_Data
 	//	*AppletUpdate_Header
 	Payload isAppletUpdate_Payload `protobuf_oneof:"Payload"`
@@ -199,9 +197,11 @@ type AppletUpdateHeader struct {
 	Signature []byte `protobuf:"bytes,1,opt,name=Signature,proto3" json:"Signature,omitempty"`
 	// Checkpoint contains a note-formatted Log Checkpoint.
 	Checkpoint []byte `protobuf:"bytes,2,opt,name=Checkpoint,proto3" json:"Checkpoint,omitempty"`
-	// Manifest is metadata about the applet, which has been logged to a firmware transparency log.
+	// Manifest is metadata about the applet, which has been logged to a firmware
+	// transparency log.
 	Manifest []byte `protobuf:"bytes,3,opt,name=Manifest,proto3" json:"Manifest,omitempty"`
-	// InclusionProof is a log inclusion proof for Manifest committed to by Checkpoint.
+	// InclusionProof is a log inclusion proof for Manifest committed to by
+	// Checkpoint.
 	InclusionProof [][]byte `protobuf:"bytes,4,rep,name=InclusionProof,proto3" json:"InclusionProof,omitempty"`
 	// LogIndex is the index of Manifest in the firmware transparency log.
 	LogIndex uint64 `protobuf:"varint,5,opt,name=LogIndex,proto3" json:"LogIndex,omitempty"`
@@ -274,13 +274,10 @@ func (x *AppletUpdateHeader) GetLogIndex() uint64 {
 	return 0
 }
 
+// Status information
 //
-//
-//Status information
-//
-//The status information format is returned on any message sent with the
-//`U2FHID_ARMORY_INF` vendor specific command.
-//
+// The status information format is returned on any message sent with the
+// `U2FHID_ARMORY_INF` vendor specific command.
 type Status struct {
 	state         protoimpl.MessageState
 	sizeCache     protoimpl.SizeCache
@@ -384,13 +381,10 @@ func (x *Status) GetWitness() *WitnessStatus {
 	return nil
 }
 
+// WitnessStatus contains witness-applet specific status information.
 //
-//
-//WitnessStatus contains witness-applet specific status information.
-//
-//This is embedded in the general Status message if the applet has provided
-//this information to the OS.
-//
+// This is embedded in the general Status message if the applet has provided
+// this information to the OS.
 type WitnessStatus struct {
 	state         protoimpl.MessageState
 	sizeCache     protoimpl.SizeCache
@@ -401,6 +395,9 @@ type WitnessStatus struct {
 	Identity string `protobuf:"bytes,1,opt,name=Identity,proto3" json:"Identity,omitempty"`
 	// IP is a string representation of the witness applet's current IP address.
 	IP string `protobuf:"bytes,2,opt,name=IP,proto3" json:"IP,omitempty"`
+	// IdentityCounter is incremented when the device is recovered and the device
+	// needs a new witness identity.
+	IdentityCounter uint32 `protobuf:"varint,3,opt,name=IdentityCounter,proto3" json:"IdentityCounter,omitempty"`
 }
 
 func (x *WitnessStatus) Reset() {
@@ -449,6 +446,13 @@ func (x *WitnessStatus) GetIP() string {
 	return ""
 }
 
+func (x *WitnessStatus) GetIdentityCounter() uint32 {
+	if x != nil {
+		return x.IdentityCounter
+	}
+	return 0
+}
+
 type Configuration struct {
 	state         protoimpl.MessageState
 	sizeCache     protoimpl.SizeCache
@@ -629,29 +633,32 @@ var file_api_proto_rawDesc = []byte{
 	0x6b, 0x12, 0x2c, 0x0a, 0x07, 0x57, 0x69, 0x74, 0x6e, 0x65, 0x73, 0x73, 0x18, 0x08, 0x20, 0x01,
 	0x28, 0x0b, 0x32, 0x12, 0x2e, 0x61, 0x70, 0x69, 0x2e, 0x57, 0x69, 0x74, 0x6e, 0x65, 0x73, 0x73,
 	0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x07, 0x57, 0x69, 0x74, 0x6e, 0x65, 0x73, 0x73, 0x22,
-	0x3b, 0x0a, 0x0d, 0x57, 0x69, 0x74, 0x6e, 0x65, 0x73, 0x73, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73,
+	0x65, 0x0a, 0x0d, 0x57, 0x69, 0x74, 0x6e, 0x65, 0x73, 0x73, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73,
 	0x12, 0x1a, 0x0a, 0x08, 0x49, 0x64, 0x65, 0x6e, 0x74, 0x69, 0x74, 0x79, 0x18, 0x01, 0x20, 0x01,
 	0x28, 0x09, 0x52, 0x08, 0x49, 0x64, 0x65, 0x6e, 0x74, 0x69, 0x74, 0x79, 0x12, 0x0e, 0x0a, 0x02,
-	0x49, 0x50, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x49, 0x50, 0x22, 0xa1, 0x01, 0x0a,
-	0x0d, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x75, 0x72, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x12,
-	0x0a, 0x04, 0x44, 0x48, 0x43, 0x50, 0x18, 0x01, 0x20, 0x01, 0x28, 0x08, 0x52, 0x04, 0x44, 0x48,
-	0x43, 0x50, 0x12, 0x0e, 0x0a, 0x02, 0x49, 0x50, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02,
-	0x49, 0x50, 0x12, 0x18, 0x0a, 0x07, 0x4e, 0x65, 0x74, 0x6d, 0x61, 0x73, 0x6b, 0x18, 0x03, 0x20,
-	0x01, 0x28, 0x09, 0x52, 0x07, 0x4e, 0x65, 0x74, 0x6d, 0x61, 0x73, 0x6b, 0x12, 0x18, 0x0a, 0x07,
-	0x47, 0x61, 0x74, 0x65, 0x77, 0x61, 0x79, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x47,
-	0x61, 0x74, 0x65, 0x77, 0x61, 0x79, 0x12, 0x1a, 0x0a, 0x08, 0x52, 0x65, 0x73, 0x6f, 0x6c, 0x76,
-	0x65, 0x72, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x52, 0x65, 0x73, 0x6f, 0x6c, 0x76,
-	0x65, 0x72, 0x12, 0x1c, 0x0a, 0x09, 0x4e, 0x54, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x18,
-	0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x4e, 0x54, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72,
-	0x22, 0x4a, 0x0a, 0x08, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x24, 0x0a, 0x05,
-	0x45, 0x72, 0x72, 0x6f, 0x72, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x0e, 0x2e, 0x61, 0x70,
-	0x69, 0x2e, 0x45, 0x72, 0x72, 0x6f, 0x72, 0x43, 0x6f, 0x64, 0x65, 0x52, 0x05, 0x45, 0x72, 0x72,
-	0x6f, 0x72, 0x12, 0x18, 0x0a, 0x07, 0x50, 0x61, 0x79, 0x6c, 0x6f, 0x61, 0x64, 0x18, 0x02, 0x20,
-	0x01, 0x28, 0x0c, 0x52, 0x07, 0x50, 0x61, 0x79, 0x6c, 0x6f, 0x61, 0x64, 0x2a, 0x28, 0x0a, 0x09,
-	0x45, 0x72, 0x72, 0x6f, 0x72, 0x43, 0x6f, 0x64, 0x65, 0x12, 0x08, 0x0a, 0x04, 0x4e, 0x4f, 0x4e,
-	0x45, 0x10, 0x00, 0x12, 0x11, 0x0a, 0x0d, 0x47, 0x45, 0x4e, 0x45, 0x52, 0x49, 0x43, 0x5f, 0x45,
-	0x52, 0x52, 0x4f, 0x52, 0x10, 0x01, 0x42, 0x08, 0x5a, 0x06, 0x2e, 0x2f, 0x3b, 0x61, 0x70, 0x69,
-	0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
+	0x49, 0x50, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x49, 0x50, 0x12, 0x28, 0x0a, 0x0f,
+	0x49, 0x64, 0x65, 0x6e, 0x74, 0x69, 0x74, 0x79, 0x43, 0x6f, 0x75, 0x6e, 0x74, 0x65, 0x72, 0x18,
+	0x03, 0x20, 0x01, 0x28, 0x0d, 0x52, 0x0f, 0x49, 0x64, 0x65, 0x6e, 0x74, 0x69, 0x74, 0x79, 0x43,
+	0x6f, 0x75, 0x6e, 0x74, 0x65, 0x72, 0x22, 0xa1, 0x01, 0x0a, 0x0d, 0x43, 0x6f, 0x6e, 0x66, 0x69,
+	0x67, 0x75, 0x72, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x12, 0x0a, 0x04, 0x44, 0x48, 0x43, 0x50,
+	0x18, 0x01, 0x20, 0x01, 0x28, 0x08, 0x52, 0x04, 0x44, 0x48, 0x43, 0x50, 0x12, 0x0e, 0x0a, 0x02,
+	0x49, 0x50, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x49, 0x50, 0x12, 0x18, 0x0a, 0x07,
+	0x4e, 0x65, 0x74, 0x6d, 0x61, 0x73, 0x6b, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x4e,
+	0x65, 0x74, 0x6d, 0x61, 0x73, 0x6b, 0x12, 0x18, 0x0a, 0x07, 0x47, 0x61, 0x74, 0x65, 0x77, 0x61,
+	0x79, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x47, 0x61, 0x74, 0x65, 0x77, 0x61, 0x79,
+	0x12, 0x1a, 0x0a, 0x08, 0x52, 0x65, 0x73, 0x6f, 0x6c, 0x76, 0x65, 0x72, 0x18, 0x05, 0x20, 0x01,
+	0x28, 0x09, 0x52, 0x08, 0x52, 0x65, 0x73, 0x6f, 0x6c, 0x76, 0x65, 0x72, 0x12, 0x1c, 0x0a, 0x09,
+	0x4e, 0x54, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52,
+	0x09, 0x4e, 0x54, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x22, 0x4a, 0x0a, 0x08, 0x52, 0x65,
+	0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x24, 0x0a, 0x05, 0x45, 0x72, 0x72, 0x6f, 0x72, 0x18,
+	0x01, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x0e, 0x2e, 0x61, 0x70, 0x69, 0x2e, 0x45, 0x72, 0x72, 0x6f,
+	0x72, 0x43, 0x6f, 0x64, 0x65, 0x52, 0x05, 0x45, 0x72, 0x72, 0x6f, 0x72, 0x12, 0x18, 0x0a, 0x07,
+	0x50, 0x61, 0x79, 0x6c, 0x6f, 0x61, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x07, 0x50,
+	0x61, 0x79, 0x6c, 0x6f, 0x61, 0x64, 0x2a, 0x28, 0x0a, 0x09, 0x45, 0x72, 0x72, 0x6f, 0x72, 0x43,
+	0x6f, 0x64, 0x65, 0x12, 0x08, 0x0a, 0x04, 0x4e, 0x4f, 0x4e, 0x45, 0x10, 0x00, 0x12, 0x11, 0x0a,
+	0x0d, 0x47, 0x45, 0x4e, 0x45, 0x52, 0x49, 0x43, 0x5f, 0x45, 0x52, 0x52, 0x4f, 0x52, 0x10, 0x01,
+	0x42, 0x08, 0x5a, 0x06, 0x2e, 0x2f, 0x3b, 0x61, 0x70, 0x69, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74,
+	0x6f, 0x33,
 }
 
 var (
diff --git a/api/api.proto b/api/api.proto
index f758e2a..471f4aa 100644
--- a/api/api.proto
+++ b/api/api.proto
@@ -16,10 +16,10 @@
 
 syntax = "proto3";
 
-option go_package = "./;api";
-
 package api;
 
+option go_package = "./;api";
+
 /*
 
 Trusted Applet updates are supported with a sequence that splits the applet
@@ -79,15 +79,14 @@ The `TotalChunks` value indicates the total number of chunks for the update,
 
 */
 message AppletUpdate {
-	uint32 Total = 1;
-	uint32 Seq   = 2;
-	oneof Payload {
-		bytes  Data  = 3;
-		AppletUpdateHeader Header = 4;
-	};
+  uint32 Total = 1;
+  uint32 Seq = 2;
+  oneof Payload {
+    bytes Data = 3;
+    AppletUpdateHeader Header = 4;
+  }
 }
 
-
 /*
 
 AppletUpdateHeader
@@ -99,16 +98,18 @@ It's sent in the zero-th `AppletUpdate` message during an OTA.
 */
 
 message AppletUpdateHeader {
-	// Signature holds the signature over the applet.
-	bytes  Signature  = 1;
-	// Checkpoint contains a note-formatted Log Checkpoint.
-	bytes  Checkpoint = 2;
-	// Manifest is metadata about the applet, which has been logged to a firmware transparency log.
-	bytes  Manifest = 3;
-	// InclusionProof is a log inclusion proof for Manifest committed to by Checkpoint.
-	repeated bytes InclusionProof = 4;
-	// LogIndex is the index of Manifest in the firmware transparency log.
-	uint64	LogIndex = 5; 
+  // Signature holds the signature over the applet.
+  bytes Signature = 1;
+  // Checkpoint contains a note-formatted Log Checkpoint.
+  bytes Checkpoint = 2;
+  // Manifest is metadata about the applet, which has been logged to a firmware
+  // transparency log.
+  bytes Manifest = 3;
+  // InclusionProof is a log inclusion proof for Manifest committed to by
+  // Checkpoint.
+  repeated bytes InclusionProof = 4;
+  // LogIndex is the index of Manifest in the firmware transparency log.
+  uint64 LogIndex = 5;
 }
 
 /*
@@ -120,14 +121,14 @@ The status information format is returned on any message sent with the
 
 */
 message Status {
-	string Serial         = 1;
-	bool   HAB            = 2;
-	string Revision       = 3;
-	string Build          = 4;
-	uint32 Version        = 5;
-	string Runtime        = 6;
-	bool   Link           = 7;
-	WitnessStatus Witness = 8;
+  string Serial = 1;
+  bool HAB = 2;
+  string Revision = 3;
+  string Build = 4;
+  uint32 Version = 5;
+  string Runtime = 6;
+  bool Link = 7;
+  WitnessStatus Witness = 8;
 }
 
 /*
@@ -139,11 +140,14 @@ this information to the OS.
 
 */
 message WitnessStatus {
-	// Identity is the note-formatted public key which can be used to verify
-	// checkpoints cosigned by this witness.
-	string Identity  = 1;
-	// IP is a string representation of the witness applet's current IP address.
-	string IP        = 2;
+  // Identity is the note-formatted public key which can be used to verify
+  // checkpoints cosigned by this witness.
+  string Identity = 1;
+  // IP is a string representation of the witness applet's current IP address.
+  string IP = 2;
+  // IdentityCounter is incremented when the device is recovered and the device
+  // needs a new witness identity.
+  uint32 IdentityCounter = 3;
 }
 
 /*
@@ -159,23 +163,23 @@ command.
 */
 
 message Configuration {
-	bool   DHCP      = 1;
-	string IP        = 2;
-	string Netmask   = 3;
-	string Gateway   = 4;
-	string Resolver  = 5;
-	string NTPServer = 6;
+  bool DHCP = 1;
+  string IP = 2;
+  string Netmask = 3;
+  string Gateway = 4;
+  string Resolver = 5;
+  string NTPServer = 6;
 }
 
 message Response {
-	ErrorCode Error   = 1;
-	bytes     Payload = 2;
+  ErrorCode Error = 1;
+  bytes Payload = 2;
 }
 
 enum ErrorCode {
-	NONE = 0;
+  NONE = 0;
 
-	// GENERIC_ERROR is returned in case of a generic error, in this case
-	// Payload might contain the error string.
-	GENERIC_ERROR = 1;
+  // GENERIC_ERROR is returned in case of a generic error, in this case
+  // Payload might contain the error string.
+  GENERIC_ERROR = 1;
 }
diff --git a/api/rpc/rpc.go b/api/rpc/rpc.go
index 19d8a97..9a4545a 100644
--- a/api/rpc/rpc.go
+++ b/api/rpc/rpc.go
@@ -49,6 +49,9 @@ type WitnessStatus struct {
 	Identity string
 	// IP is the currently-assigned IP address of the witness applet.
 	IP string
+	// IdentityCounter is incremented when the device is recovered and the device
+	// needs a new witness identity.
+	IdentityCounter uint32
 }
 
 // FirmwareUpdate represents a firmware update.

From f4c93242de37c55f89ade2de534906a78100a1c0 Mon Sep 17 00:00:00 2001
From: Jay <jayhou@google.com>
Date: Tue, 30 Jan 2024 11:58:22 +0000
Subject: [PATCH 10/16] Move identity counter out of WitnessStatus.

---
 api/api.go        |  2 +-
 api/api.pb.go     | 38 +++++++++++++++++++-------------------
 api/api.proto     |  6 +++---
 api/rpc/rpc.go    |  3 ---
 trusted_os/ctl.go |  1 +
 5 files changed, 24 insertions(+), 26 deletions(-)

diff --git a/api/api.go b/api/api.go
index fe71a31..cfd959d 100644
--- a/api/api.go
+++ b/api/api.go
@@ -98,10 +98,10 @@ func (p *Status) Print() string {
 	status.WriteString(fmt.Sprintf("Version ....................: %d (%s)\n", p.Version, time.Unix(int64(p.Version), 0)))
 	status.WriteString(fmt.Sprintf("Runtime ....................: %s\n", p.Runtime))
 	status.WriteString(fmt.Sprintf("Link .......................: %v\n", p.Link))
+	status.WriteString(fmt.Sprintf("IdentityCounter ............: %d", p.IdentityCounter))
 	if p.Witness != nil {
 		status.WriteString(fmt.Sprintf("Witness/Identity ...........: %v\n", p.Witness.Identity))
 		status.WriteString(fmt.Sprintf("Witness/IP .................: %v", p.Witness.IP))
-		status.WriteString(fmt.Sprintf("Witness/IdentityCounter ....: %d", p.Witness.IdentityCounter))
 	} else {
 		status.WriteString(fmt.Sprint("Witness ....................: <no status>"))
 	}
diff --git a/api/api.pb.go b/api/api.pb.go
index 0aed4b4..b106ed1 100644
--- a/api/api.pb.go
+++ b/api/api.pb.go
@@ -291,6 +291,9 @@ type Status struct {
 	Runtime  string         `protobuf:"bytes,6,opt,name=Runtime,proto3" json:"Runtime,omitempty"`
 	Link     bool           `protobuf:"varint,7,opt,name=Link,proto3" json:"Link,omitempty"`
 	Witness  *WitnessStatus `protobuf:"bytes,8,opt,name=Witness,proto3" json:"Witness,omitempty"`
+	// IdentityCounter is incremented when the device is recovered and the device
+	// needs a new witness identity.
+	IdentityCounter uint32 `protobuf:"varint,9,opt,name=IdentityCounter,proto3" json:"IdentityCounter,omitempty"`
 }
 
 func (x *Status) Reset() {
@@ -381,6 +384,13 @@ func (x *Status) GetWitness() *WitnessStatus {
 	return nil
 }
 
+func (x *Status) GetIdentityCounter() uint32 {
+	if x != nil {
+		return x.IdentityCounter
+	}
+	return 0
+}
+
 // WitnessStatus contains witness-applet specific status information.
 //
 // This is embedded in the general Status message if the applet has provided
@@ -395,9 +405,6 @@ type WitnessStatus struct {
 	Identity string `protobuf:"bytes,1,opt,name=Identity,proto3" json:"Identity,omitempty"`
 	// IP is a string representation of the witness applet's current IP address.
 	IP string `protobuf:"bytes,2,opt,name=IP,proto3" json:"IP,omitempty"`
-	// IdentityCounter is incremented when the device is recovered and the device
-	// needs a new witness identity.
-	IdentityCounter uint32 `protobuf:"varint,3,opt,name=IdentityCounter,proto3" json:"IdentityCounter,omitempty"`
 }
 
 func (x *WitnessStatus) Reset() {
@@ -446,13 +453,6 @@ func (x *WitnessStatus) GetIP() string {
 	return ""
 }
 
-func (x *WitnessStatus) GetIdentityCounter() uint32 {
-	if x != nil {
-		return x.IdentityCounter
-	}
-	return 0
-}
-
 type Configuration struct {
 	state         protoimpl.MessageState
 	sizeCache     protoimpl.SizeCache
@@ -619,7 +619,7 @@ var file_api_proto_rawDesc = []byte{
 	0x18, 0x04, 0x20, 0x03, 0x28, 0x0c, 0x52, 0x0e, 0x49, 0x6e, 0x63, 0x6c, 0x75, 0x73, 0x69, 0x6f,
 	0x6e, 0x50, 0x72, 0x6f, 0x6f, 0x66, 0x12, 0x1a, 0x0a, 0x08, 0x4c, 0x6f, 0x67, 0x49, 0x6e, 0x64,
 	0x65, 0x78, 0x18, 0x05, 0x20, 0x01, 0x28, 0x04, 0x52, 0x08, 0x4c, 0x6f, 0x67, 0x49, 0x6e, 0x64,
-	0x65, 0x78, 0x22, 0xda, 0x01, 0x0a, 0x06, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x12, 0x16, 0x0a,
+	0x65, 0x78, 0x22, 0x84, 0x02, 0x0a, 0x06, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x12, 0x16, 0x0a,
 	0x06, 0x53, 0x65, 0x72, 0x69, 0x61, 0x6c, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x53,
 	0x65, 0x72, 0x69, 0x61, 0x6c, 0x12, 0x10, 0x0a, 0x03, 0x48, 0x41, 0x42, 0x18, 0x02, 0x20, 0x01,
 	0x28, 0x08, 0x52, 0x03, 0x48, 0x41, 0x42, 0x12, 0x1a, 0x0a, 0x08, 0x52, 0x65, 0x76, 0x69, 0x73,
@@ -632,14 +632,14 @@ var file_api_proto_rawDesc = []byte{
 	0x04, 0x4c, 0x69, 0x6e, 0x6b, 0x18, 0x07, 0x20, 0x01, 0x28, 0x08, 0x52, 0x04, 0x4c, 0x69, 0x6e,
 	0x6b, 0x12, 0x2c, 0x0a, 0x07, 0x57, 0x69, 0x74, 0x6e, 0x65, 0x73, 0x73, 0x18, 0x08, 0x20, 0x01,
 	0x28, 0x0b, 0x32, 0x12, 0x2e, 0x61, 0x70, 0x69, 0x2e, 0x57, 0x69, 0x74, 0x6e, 0x65, 0x73, 0x73,
-	0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x07, 0x57, 0x69, 0x74, 0x6e, 0x65, 0x73, 0x73, 0x22,
-	0x65, 0x0a, 0x0d, 0x57, 0x69, 0x74, 0x6e, 0x65, 0x73, 0x73, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73,
-	0x12, 0x1a, 0x0a, 0x08, 0x49, 0x64, 0x65, 0x6e, 0x74, 0x69, 0x74, 0x79, 0x18, 0x01, 0x20, 0x01,
-	0x28, 0x09, 0x52, 0x08, 0x49, 0x64, 0x65, 0x6e, 0x74, 0x69, 0x74, 0x79, 0x12, 0x0e, 0x0a, 0x02,
-	0x49, 0x50, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x49, 0x50, 0x12, 0x28, 0x0a, 0x0f,
-	0x49, 0x64, 0x65, 0x6e, 0x74, 0x69, 0x74, 0x79, 0x43, 0x6f, 0x75, 0x6e, 0x74, 0x65, 0x72, 0x18,
-	0x03, 0x20, 0x01, 0x28, 0x0d, 0x52, 0x0f, 0x49, 0x64, 0x65, 0x6e, 0x74, 0x69, 0x74, 0x79, 0x43,
-	0x6f, 0x75, 0x6e, 0x74, 0x65, 0x72, 0x22, 0xa1, 0x01, 0x0a, 0x0d, 0x43, 0x6f, 0x6e, 0x66, 0x69,
+	0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x07, 0x57, 0x69, 0x74, 0x6e, 0x65, 0x73, 0x73, 0x12,
+	0x28, 0x0a, 0x0f, 0x49, 0x64, 0x65, 0x6e, 0x74, 0x69, 0x74, 0x79, 0x43, 0x6f, 0x75, 0x6e, 0x74,
+	0x65, 0x72, 0x18, 0x09, 0x20, 0x01, 0x28, 0x0d, 0x52, 0x0f, 0x49, 0x64, 0x65, 0x6e, 0x74, 0x69,
+	0x74, 0x79, 0x43, 0x6f, 0x75, 0x6e, 0x74, 0x65, 0x72, 0x22, 0x3b, 0x0a, 0x0d, 0x57, 0x69, 0x74,
+	0x6e, 0x65, 0x73, 0x73, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x12, 0x1a, 0x0a, 0x08, 0x49, 0x64,
+	0x65, 0x6e, 0x74, 0x69, 0x74, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x49, 0x64,
+	0x65, 0x6e, 0x74, 0x69, 0x74, 0x79, 0x12, 0x0e, 0x0a, 0x02, 0x49, 0x50, 0x18, 0x02, 0x20, 0x01,
+	0x28, 0x09, 0x52, 0x02, 0x49, 0x50, 0x22, 0xa1, 0x01, 0x0a, 0x0d, 0x43, 0x6f, 0x6e, 0x66, 0x69,
 	0x67, 0x75, 0x72, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x12, 0x0a, 0x04, 0x44, 0x48, 0x43, 0x50,
 	0x18, 0x01, 0x20, 0x01, 0x28, 0x08, 0x52, 0x04, 0x44, 0x48, 0x43, 0x50, 0x12, 0x0e, 0x0a, 0x02,
 	0x49, 0x50, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x49, 0x50, 0x12, 0x18, 0x0a, 0x07,
diff --git a/api/api.proto b/api/api.proto
index 471f4aa..d98b383 100644
--- a/api/api.proto
+++ b/api/api.proto
@@ -129,6 +129,9 @@ message Status {
   string Runtime = 6;
   bool Link = 7;
   WitnessStatus Witness = 8;
+  // IdentityCounter is incremented when the device is recovered and the device
+  // needs a new witness identity.
+  uint32 IdentityCounter = 9;
 }
 
 /*
@@ -145,9 +148,6 @@ message WitnessStatus {
   string Identity = 1;
   // IP is a string representation of the witness applet's current IP address.
   string IP = 2;
-  // IdentityCounter is incremented when the device is recovered and the device
-  // needs a new witness identity.
-  uint32 IdentityCounter = 3;
 }
 
 /*
diff --git a/api/rpc/rpc.go b/api/rpc/rpc.go
index 9a4545a..19d8a97 100644
--- a/api/rpc/rpc.go
+++ b/api/rpc/rpc.go
@@ -49,9 +49,6 @@ type WitnessStatus struct {
 	Identity string
 	// IP is the currently-assigned IP address of the witness applet.
 	IP string
-	// IdentityCounter is incremented when the device is recovered and the device
-	// needs a new witness identity.
-	IdentityCounter uint32
 }
 
 // FirmwareUpdate represents a firmware update.
diff --git a/trusted_os/ctl.go b/trusted_os/ctl.go
index 5fa8f7e..8ff11a5 100644
--- a/trusted_os/ctl.go
+++ b/trusted_os/ctl.go
@@ -61,6 +61,7 @@ func getStatus() (s *api.Status) {
 		Build:    Build,
 		Version:  version,
 		Runtime:  fmt.Sprintf("%s %s/%s", runtime.Version(), runtime.GOOS, runtime.GOARCH),
+		// TODO(jayhou): set IdentityCounter here.
 	}
 	if witnessStatus != nil {
 		s.Witness = &api.WitnessStatus{

From 350721de7701bda289c8dc7054971ccc64a59654 Mon Sep 17 00:00:00 2001
From: Jay <jayhou@google.com>
Date: Tue, 30 Jan 2024 12:17:32 +0000
Subject: [PATCH 11/16] Set identity counter in status.

---
 trusted_os/ctl.go       |  8 +++++++-
 trusted_os/main.go      |  3 ++-
 trusted_os/rpmb.go      | 13 +++++++++++++
 trusted_os/rpmb_fake.go |  9 +++++++++
 4 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/trusted_os/ctl.go b/trusted_os/ctl.go
index 8ff11a5..65463c2 100644
--- a/trusted_os/ctl.go
+++ b/trusted_os/ctl.go
@@ -46,6 +46,7 @@ type controlInterface struct {
 	sync.Mutex
 
 	Device *usb.Device
+	rpmb   *RPMB
 	RPC    *RPC
 
 	ota *otaBuffer
@@ -61,7 +62,12 @@ func getStatus() (s *api.Status) {
 		Build:    Build,
 		Version:  version,
 		Runtime:  fmt.Sprintf("%s %s/%s", runtime.Version(), runtime.GOOS, runtime.GOARCH),
-		// TODO(jayhou): set IdentityCounter here.
+	}
+	count, err := rpmb.witnessIdentity()
+	if err != nil {
+		log.Errorf("cannot get witness identity counter: %v", err)
+	} else {
+		s.WitnessIdentity = count
 	}
 	if witnessStatus != nil {
 		s.Witness = &api.WitnessStatus{
diff --git a/trusted_os/main.go b/trusted_os/main.go
index b3d8bbc..ded0e5d 100644
--- a/trusted_os/main.go
+++ b/trusted_os/main.go
@@ -153,7 +153,8 @@ func main() {
 	}
 
 	ctl := &controlInterface{
-		RPC: rpc,
+		RPC:  rpc,
+		RPMB: rpmb,
 	}
 
 	// TODO: disable for now
diff --git a/trusted_os/rpmb.go b/trusted_os/rpmb.go
index 026205f..4f554dc 100644
--- a/trusted_os/rpmb.go
+++ b/trusted_os/rpmb.go
@@ -219,6 +219,19 @@ func (r *RPMB) transfer(offset uint16, buf []byte, n *uint32, write bool) (err e
 	return
 }
 
+// witnessIdentity gets the witness identity counter from the RPMB area.
+func (r *RPMB) witnessIdentity() (counter uint32, err error) {
+	if r.partition == nil {
+		return 0, errors.New("RPMB has not been initialized")
+	}
+
+	rBuf := make([]byte, witnessIdentityCounterLength)
+	if err = r.partition.Read(witnessIdentityCounter, rBuf); err != nil {
+		return 0, err
+	}
+	return binary.BigEndian.Uint32(rBuf), nil
+}
+
 // incrementWitnessIdentity increments the counter in the RPMB area to
 // differentiate a new witness identity.
 func (r *RPMB) incrementWitnessIdentity() (err error) {
diff --git a/trusted_os/rpmb_fake.go b/trusted_os/rpmb_fake.go
index 4d6f22a..0ddd502 100644
--- a/trusted_os/rpmb_fake.go
+++ b/trusted_os/rpmb_fake.go
@@ -135,6 +135,15 @@ func (r *RPMB) transfer(sector uint16, buf []byte, n *uint32, write bool) (err e
 	return
 }
 
+// witnessIdentity gets the witness identity counter from the RPMB area.
+func (r *RPMB) witnessIdentity() (counter uint32, err error) {
+	rBuf := make([]byte, witnessIdentityCounterLength)
+	if err := r.transfer(witnessIdentityCounterSector, rBuf, nil, false); err != nil {
+		return err
+	}
+	return binary.BigEndian.Uint32(rBuf)
+}
+
 // incrementWitnessIdentity increments the counter in the RPMB area to
 // differentiate a new witness identity.
 func (r *RPMB) incrementWitnessIdentity() error {

From ad8e820f3f3fdb7db792635471123e70e70814ee Mon Sep 17 00:00:00 2001
From: Jay <jayhou@google.com>
Date: Tue, 30 Jan 2024 12:46:59 +0000
Subject: [PATCH 12/16] Fix getStatus.

---
 trusted_os/ctl.go  | 17 ++++++++++-------
 trusted_os/main.go |  3 +--
 trusted_os/rpc.go  |  3 ++-
 3 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/trusted_os/ctl.go b/trusted_os/ctl.go
index 65463c2..56f377c 100644
--- a/trusted_os/ctl.go
+++ b/trusted_os/ctl.go
@@ -46,13 +46,12 @@ type controlInterface struct {
 	sync.Mutex
 
 	Device *usb.Device
-	rpmb   *RPMB
 	RPC    *RPC
 
 	ota *otaBuffer
 }
 
-func getStatus() (s *api.Status) {
+func (ctl *controlInterface) getStatus(rpmb *RPMB) (s *api.Status) {
 	version, _ := parseVersion(Version)
 
 	s = &api.Status{
@@ -63,11 +62,15 @@ func getStatus() (s *api.Status) {
 		Version:  version,
 		Runtime:  fmt.Sprintf("%s %s/%s", runtime.Version(), runtime.GOOS, runtime.GOARCH),
 	}
-	count, err := rpmb.witnessIdentity()
-	if err != nil {
-		log.Errorf("cannot get witness identity counter: %v", err)
+	if rpmb == nil {
+		log.Printf("cannot get witness identity counter because no RPMB passed in to controlInterface")
 	} else {
-		s.WitnessIdentity = count
+		count, err := rpmb.witnessIdentity()
+		if err != nil {
+			log.Printf("cannot get witness identity counter: %v", err)
+		} else {
+			s.IdentityCounter = count
+		}
 	}
 	if witnessStatus != nil {
 		s.Witness = &api.WitnessStatus{
@@ -93,7 +96,7 @@ func (ctl *controlInterface) HandleMessage(_ []byte) (_ []byte) {
 }
 
 func (ctl *controlInterface) Status(_ []byte) (res []byte) {
-	res, _ = proto.Marshal(getStatus())
+	res, _ = proto.Marshal(ctl.getStatus(nil))
 	return
 }
 
diff --git a/trusted_os/main.go b/trusted_os/main.go
index ded0e5d..b3d8bbc 100644
--- a/trusted_os/main.go
+++ b/trusted_os/main.go
@@ -153,8 +153,7 @@ func main() {
 	}
 
 	ctl := &controlInterface{
-		RPC:  rpc,
-		RPMB: rpmb,
+		RPC: rpc,
 	}
 
 	// TODO: disable for now
diff --git a/trusted_os/rpc.go b/trusted_os/rpc.go
index 616b630..c2f15dd 100644
--- a/trusted_os/rpc.go
+++ b/trusted_os/rpc.go
@@ -41,6 +41,7 @@ type RPC struct {
 	Storage     Card
 	Ctx         *monitor.ExecCtx
 	Cfg         []byte
+	Ctl         controlInterface
 	Diversifier [32]byte
 }
 
@@ -105,7 +106,7 @@ func (r *RPC) Status(_ any, status *api.Status) error {
 		return errors.New("invalid argument")
 	}
 
-	s := getStatus()
+	s := r.Ctl.getStatus(r.RPMB)
 	*status = *s
 
 	return nil

From 51555b4e984749eb9ec470c35ac64ba53c9027b2 Mon Sep 17 00:00:00 2001
From: Jay <jayhou@google.com>
Date: Thu, 1 Feb 2024 10:44:20 +0000
Subject: [PATCH 13/16] Don't need to pass in RPMB to ctl.

---
 trusted_os/ctl.go | 11 ++++++-----
 trusted_os/rpc.go |  2 +-
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/trusted_os/ctl.go b/trusted_os/ctl.go
index 4188416..9ad50c2 100644
--- a/trusted_os/ctl.go
+++ b/trusted_os/ctl.go
@@ -51,7 +51,7 @@ type controlInterface struct {
 	ota *otaBuffer
 }
 
-func (ctl *controlInterface) getStatus(rpmb *RPMB) (s *api.Status) {
+func (ctl *controlInterface) getStatus() (s *api.Status) {
 	version, _ := parseVersion(Version)
 
 	s = &api.Status{
@@ -61,10 +61,10 @@ func (ctl *controlInterface) getStatus(rpmb *RPMB) (s *api.Status) {
 		Build:    Build,
 		Version:  version,
 		Runtime:  fmt.Sprintf("%s %s/%s", runtime.Version(), runtime.GOOS, runtime.GOARCH),
-		// TODO(jayhou): set IdentityCounter here.
 	}
-	if rpmb == nil {
-		log.Printf("cannot get witness identity counter because no RPMB passed in to controlInterface")
+
+	if rpmb := ctl.RPC.RPMB; rpmb == nil {
+		log.Printf("cannot get witness identity counter because RPMB is nil")
 	} else {
 		count, err := rpmb.witnessIdentity()
 		if err != nil {
@@ -73,6 +73,7 @@ func (ctl *controlInterface) getStatus(rpmb *RPMB) (s *api.Status) {
 			s.IdentityCounter = count
 		}
 	}
+
 	if witnessStatus != nil {
 		s.Witness = &api.WitnessStatus{
 			Identity: witnessStatus.Identity,
@@ -97,7 +98,7 @@ func (ctl *controlInterface) HandleMessage(_ []byte) (_ []byte) {
 }
 
 func (ctl *controlInterface) Status(_ []byte) (res []byte) {
-	res, _ = proto.Marshal(ctl.getStatus(nil))
+	res, _ = proto.Marshal(ctl.getStatus())
 	return
 }
 
diff --git a/trusted_os/rpc.go b/trusted_os/rpc.go
index c2f15dd..284700f 100644
--- a/trusted_os/rpc.go
+++ b/trusted_os/rpc.go
@@ -106,7 +106,7 @@ func (r *RPC) Status(_ any, status *api.Status) error {
 		return errors.New("invalid argument")
 	}
 
-	s := r.Ctl.getStatus(r.RPMB)
+	s := r.Ctl.getStatus()
 	*status = *s
 
 	return nil

From 187befacdcb7b0a514c15b712ac94e22832f6760 Mon Sep 17 00:00:00 2001
From: Jay <jayhou@google.com>
Date: Thu, 1 Feb 2024 20:01:23 +0000
Subject: [PATCH 14/16] Fix compilation errors.

---
 trusted_os/ctl.go       |  5 +++--
 trusted_os/rpc.go       |  2 +-
 trusted_os/rpmb.go      | 18 +++++++++---------
 trusted_os/rpmb_fake.go | 21 +++++++++++++--------
 4 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/trusted_os/ctl.go b/trusted_os/ctl.go
index 9ad50c2..25c39af 100644
--- a/trusted_os/ctl.go
+++ b/trusted_os/ctl.go
@@ -63,10 +63,11 @@ func (ctl *controlInterface) getStatus() (s *api.Status) {
 		Runtime:  fmt.Sprintf("%s %s/%s", runtime.Version(), runtime.GOOS, runtime.GOARCH),
 	}
 
-	if rpmb := ctl.RPC.RPMB; rpmb == nil {
+	rpmb := ctl.RPC.RPMB
+	if rpmb == nil {
 		log.Printf("cannot get witness identity counter because RPMB is nil")
 	} else {
-		count, err := rpmb.witnessIdentity()
+		count, err := rpmb.witnessIdentityCounter()
 		if err != nil {
 			log.Printf("cannot get witness identity counter: %v", err)
 		} else {
diff --git a/trusted_os/rpc.go b/trusted_os/rpc.go
index 284700f..62bad47 100644
--- a/trusted_os/rpc.go
+++ b/trusted_os/rpc.go
@@ -178,7 +178,7 @@ func (r *RPC) ReadRPMB(buf []byte, n *uint32) error {
 // value stored in this area.
 func (r *RPC) ReadIdentityCounterRPMB(_ any, counter *uint32) error {
 	buf := make([]byte, witnessIdentityCounterLength)
-	if err := r.RPMB.transfer(witnessIdentityCounter, buf, nil, false); err != nil {
+	if err := r.RPMB.transfer(witnessIdentityCounterSector, buf, nil, false); err != nil {
 		return err
 	}
 
diff --git a/trusted_os/rpmb.go b/trusted_os/rpmb.go
index 4f554dc..af36266 100644
--- a/trusted_os/rpmb.go
+++ b/trusted_os/rpmb.go
@@ -57,8 +57,8 @@ const (
 
 	// witness identity counter length - uint32
 	witnessIdentityCounterLength = 4
-	// RPMB witness identity counter
-	witnessIdentityCounter = 7
+	// RPMB sector for witness identity counter
+	witnessIdentityCounterSector = 7
 
 	diversifierMAC = "ArmoryWitnessMAC"
 	iter           = 4096
@@ -219,29 +219,29 @@ func (r *RPMB) transfer(offset uint16, buf []byte, n *uint32, write bool) (err e
 	return
 }
 
-// witnessIdentity gets the witness identity counter from the RPMB area.
-func (r *RPMB) witnessIdentity() (counter uint32, err error) {
+// witnessIdentityCounter gets the witness identity counter from the RPMB area.
+func (r *RPMB) witnessIdentityCounter() (counter uint32, err error) {
 	if r.partition == nil {
 		return 0, errors.New("RPMB has not been initialized")
 	}
 
 	rBuf := make([]byte, witnessIdentityCounterLength)
-	if err = r.partition.Read(witnessIdentityCounter, rBuf); err != nil {
+	if err = r.partition.Read(witnessIdentityCounterSector, rBuf); err != nil {
 		return 0, err
 	}
 	return binary.BigEndian.Uint32(rBuf), nil
 }
 
-// incrementWitnessIdentity increments the counter in the RPMB area to
+// incrementWitnessIdentityCounter increments the counter in the RPMB area to
 // differentiate a new witness identity.
-func (r *RPMB) incrementWitnessIdentity() (err error) {
+func (r *RPMB) incrementWitnessIdentityCounter() (err error) {
 	if r.partition == nil {
 		return errors.New("RPMB has not been initialized")
 	}
 
 	// Read
 	rBuf := make([]byte, witnessIdentityCounterLength)
-	if err = r.partition.Read(witnessIdentityCounter, rBuf); err != nil {
+	if err = r.partition.Read(witnessIdentityCounterSector, rBuf); err != nil {
 		return err
 	}
 	counter := binary.BigEndian.Uint32(rBuf)
@@ -252,5 +252,5 @@ func (r *RPMB) incrementWitnessIdentity() (err error) {
 	// Write
 	wBuf := make([]byte, witnessIdentityCounterLength)
 	binary.BigEndian.PutUint32(wBuf, counter)
-	return r.partition.Write(witnessIdentityCounter, wBuf)
+	return r.partition.Write(witnessIdentityCounterSector, wBuf)
 }
diff --git a/trusted_os/rpmb_fake.go b/trusted_os/rpmb_fake.go
index 0ddd502..3e3fdfb 100644
--- a/trusted_os/rpmb_fake.go
+++ b/trusted_os/rpmb_fake.go
@@ -36,6 +36,7 @@ const (
 	// RPMB sector for TA use
 	taUserSector = 3
 
+	// witness identity counter length - uint32
 	witnessIdentityCounterLength = 4
 	// RPMB sector for witness identity counter
 	witnessIdentityCounterSector = 4
@@ -45,12 +46,16 @@ const (
 )
 
 type RPMB struct {
-	mem     [numSectors][sectorLength]byte
+	mem     [numSectors][]byte
 	counter uint32
 }
 
 func newRPMB(_ Card) (*RPMB, error) {
-	return &RPMB{}, nil
+	r := &RPMB{}
+	for i := 0; i < numSectors; i++ {
+		r.mem[i] = make([]byte, sectorLength)
+	}
+	return r, nil
 }
 
 func parseVersion(s string) (version uint32, err error) {
@@ -135,18 +140,18 @@ func (r *RPMB) transfer(sector uint16, buf []byte, n *uint32, write bool) (err e
 	return
 }
 
-// witnessIdentity gets the witness identity counter from the RPMB area.
-func (r *RPMB) witnessIdentity() (counter uint32, err error) {
+// witnessIdentityCounter gets the witness identity counter from the RPMB area.
+func (r *RPMB) witnessIdentityCounter() (counter uint32, err error) {
 	rBuf := make([]byte, witnessIdentityCounterLength)
 	if err := r.transfer(witnessIdentityCounterSector, rBuf, nil, false); err != nil {
-		return err
+		return 0, err
 	}
-	return binary.BigEndian.Uint32(rBuf)
+	return binary.BigEndian.Uint32(rBuf), nil
 }
 
-// incrementWitnessIdentity increments the counter in the RPMB area to
+// incrementWitnessIdentityCounter increments the counter in the RPMB area to
 // differentiate a new witness identity.
-func (r *RPMB) incrementWitnessIdentity() error {
+func (r *RPMB) incrementWitnessIdentityCounter() error {
 	// Read
 	rBuf := make([]byte, witnessIdentityCounterLength)
 	if err := r.transfer(witnessIdentityCounterSector, rBuf, nil, false); err != nil {

From 82397d74492b48c101013743b7ec5c7eb40036ba Mon Sep 17 00:00:00 2001
From: Jay <jayhou@google.com>
Date: Fri, 2 Feb 2024 17:37:00 +0000
Subject: [PATCH 15/16] Add env vars for toggling fake storage to Dockerfile.

---
 Dockerfile | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index 0bba0e8..808f6e1 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -8,6 +8,8 @@ ARG OS_PUBLIC_KEY1
 ARG OS_PUBLIC_KEY2
 ARG GIT_SEMVER_TAG
 ARG DEBUG
+ARG FAKE_RPMB=1
+ARG FAKE_STORAGE=0
 
 # Install dependencies.
 RUN apt-get update && apt-get install -y make wget u-boot-tools binutils-arm-none-eabi
@@ -34,6 +36,8 @@ ENV LOG_ORIGIN=${LOG_ORIGIN} \
     OS_PUBLIC_KEY1="/tmp/os1.pub" \
     OS_PUBLIC_KEY2="/tmp/os2.pub" \
     GIT_SEMVER_TAG=${GIT_SEMVER_TAG} \
-    DEBUG=${DEBUG}
+    DEBUG=${DEBUG} \
+    FAKE_RPMB=${FAKE_RPMB} \
+    FAKE_STORAGE=${FAKE_STORAGE}
 
 RUN make trusted_os_release

From 2eef2ee271bac4c5aba43a6103fa286b4d587548 Mon Sep 17 00:00:00 2001
From: Jay <jayhou@google.com>
Date: Mon, 12 Feb 2024 14:48:22 +0000
Subject: [PATCH 16/16] Add back / fix fake rpmb init()

---
 trusted_os/rpmb_fake.go | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/trusted_os/rpmb_fake.go b/trusted_os/rpmb_fake.go
index 144df64..25099d0 100644
--- a/trusted_os/rpmb_fake.go
+++ b/trusted_os/rpmb_fake.go
@@ -58,6 +58,10 @@ func newRPMB(_ Card) (*RPMB, error) {
 	return r, nil
 }
 
+func (r *RPMB) init() error {
+	return nil
+}
+
 func parseVersion(s string) (version uint32, err error) {
 	v, err := strconv.Atoi(s)
 	if err != nil {