feat: add marshal/unmarshall to command/response

tpm2/marshalling.go

@@ -2,6 +2,7 @@ package tpm2
 
 import (
 	"bytes"
+	"encoding/binary"
 	"fmt"
 	"reflect"
 )
@@ -126,3 +127,251 @@ func (b *boxed[T]) unmarshal(buf *bytes.Buffer) error {
 	b.Contents = new(T)
 	return unmarshal(buf, reflect.ValueOf(b.Contents))
 }
+
+// MarshalCommand marshals a TPM command into its raw cpHash preimage format.
+// The returned bytes can be directly hashed to compute cpHash.
+//
+// Example:
+//
+//	cmdData, _ := MarshalCommand(myCmd)
+//	cpHash := sha256.Sum256(cmdData)
+//
+// Note: Encrypted command parameters (via sessions) are not currently supported.
+// The marshaled parameters are in their unencrypted form.
+func MarshalCommand[C Command[R, *R], R any](cmd C) ([]byte, error) {
+	cc := cmd.Command()
+
+	names, err := cmdNames(cmd)
+	if err != nil {
+		return nil, err
+	}
+
+	params, err := cmdParameters(cmd, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	// Build raw cpHash preimage: CommandCode {∥ Name1 {∥ Name2 {∥ Name3 }}} {∥ Parameters }
+	// See section 16.7 of TPM 2.0 specification, part 1.
+	buf := new(bytes.Buffer)
+
+	if err := binary.Write(buf, binary.BigEndian, cc); err != nil {
+		return nil, fmt.Errorf("marshalling command code: %w", err)
+	}
+
+	for i, name := range names {
+		if _, err := buf.Write(name.Buffer); err != nil {
+			return nil, fmt.Errorf("marshalling name %d: %w", i, err)
+		}
+	}
+
+	if _, err := buf.Write(params); err != nil {
+		return nil, fmt.Errorf("marshalling parameters: %w", err)
+	}
+
+	return buf.Bytes(), nil
+}
+
+// UnmarshalCommand unmarshals a raw cpHash preimage back into a TPM command.
+// The data should be the output from [MarshalCommand].
+//
+// Example:
+//
+//	cmdData, _ := MarshalCommand(myCmd)
+//	cmd, _ := UnmarshalCommand[MyCommandType](cmdData)
+//
+// Notes:
+//   - command produced from this function is not meant to be executed directly on a TPM,
+//     instead it is expected to be used for purposes such as auditing or inspection.
+//   - encrypted command parameters (via sessions) are not currently supported.
+func UnmarshalCommand[C Command[R, *R], R any](data []byte) (C, error) {
+	var cmd C
+
+	if data == nil {
+		return cmd, fmt.Errorf("data cannot be nil")
+	}
+
+	buf := bytes.NewBuffer(data)
+
+	var cc TPMCC
+	if err := binary.Read(buf, binary.BigEndian, &cc); err != nil {
+		return cmd, fmt.Errorf("unmarshalling command code: %w", err)
+	}
+
+	if cc != cmd.Command() {
+		return cmd, fmt.Errorf("command code mismatch: expected %v, got %v", cmd.Command(), cc)
+	}
+
+	expectedNames, err := cmdNames(cmd)
+	if err != nil {
+		return cmd, fmt.Errorf("getting expected names count: %w", err)
+	}
+	numNames := len(expectedNames)
+
+	names := make([]TPM2BName, numNames)
+	for i := range numNames {
+		remaining := buf.Bytes()
+		if len(remaining) == 0 {
+			return cmd, fmt.Errorf("unexpected end of data while parsing name %d", i)
+		}
+
+		nameSize, err := parseNameSize(remaining)
+		if err != nil {
+			return cmd, fmt.Errorf("parsing name %d size: %w", i, err)
+		}
+
+		if len(remaining) < nameSize {
+			return cmd, fmt.Errorf("insufficient data for name %d: need %d bytes, have %d", i, nameSize, len(remaining))
+		}
+
+		nameBytes := make([]byte, nameSize)
+		if _, err := buf.Read(nameBytes); err != nil {
+			return cmd, fmt.Errorf("reading name %d: %w", i, err)
+		}
+
+		names[i] = TPM2BName{Buffer: nameBytes}
+	}
+
+	// Populate the command's handle fields from the names
+	if err := populateHandlesFromNames(&cmd, names); err != nil {
+		return cmd, err
+	}
+
+	params := buf.Bytes()
+
+	paramsBuf := bytes.NewBuffer(params)
+	if err := unmarshalCmdParameters(paramsBuf, &cmd, nil); err != nil {
+		return cmd, err
+	}
+	return cmd, nil
+}
+
+// parseNameSize determines the size of a TPM2BName by inspecting its first bytes.
+// Returns the total size in bytes for the name.
+//
+// Case 1: Handle-based names (4 bytes)
+//   - 0x0000... → PCR
+//   - 0x02...   → HMAC Session
+//   - 0x03...   → Policy Session
+//   - 0x40...   → Permanent Values
+//
+// Case 2: Hash-based names (2 + hash_size bytes) - for all other entities
+//   - Format: nameAlg (2 bytes) || H_nameAlg (hash digest)
+//
+// See section 14 of TPM 2.0 specification, part 1.
+func parseNameSize(buf []byte) (int, error) {
+	if len(buf) < 2 {
+		return 0, fmt.Errorf("buffer too short to parse name")
+	}
+
+	firstByte := TPMHT(buf[0])
+	firstTwoBytes := binary.BigEndian.Uint16(buf[0:2])
+
+	// Case 1: Handle-based names (4 bytes)
+	switch {
+	case firstTwoBytes == 0x0000:
+		// PCR handles (pattern: 0x0000XXXX)
+		// Must check both bytes to distinguish from hash algorithms
+		// that also start with 0x00 (e.g., TPMAlgSHA256 = 0x000B)
+		return 4, nil
+	case firstByte == TPMHTHMACSession: // 0x02
+		return 4, nil
+	case firstByte == TPMHTPolicySession: // 0x03
+		return 4, nil
+	case firstByte == TPMHTPermanent: // 0x40
+		return 4, nil
+	}
+
+	// Case 2: Hash-based names (nameAlg || hash)
+	// firstTwoBytes is the algorithm ID (0x0001 to 0x00B3)
+	algID := TPMIAlgHash(firstTwoBytes)
+	hashAlg, err := algID.Hash()
+	if err != nil {
+		return 0, fmt.Errorf("unsupported hash algorithm 0x%x in name: %w", firstTwoBytes, err)
+	}
+
+	// 2 bytes for algID + hash size
+	return 2 + hashAlg.Size(), nil
+}
+
+// MarshalResponse marshals a TPM response into its raw rpHash preimage format.
+// The returned bytes can be directly hashed to compute rpHash.
+//
+// Example:
+//
+//	rspData, _ := MarshalResponse(myCmd, myRsp)
+//	rpHash := sha256.Sum256(rspData)
+//
+// Note: Encrypted response parameters (via sessions) are not currently supported.
+func MarshalResponse[C Command[R, *R], R any](cmd C, rsp *R) ([]byte, error) {
+	cc := cmd.Command()
+
+	params, err := marshalRspParameters(rsp, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	// Build raw rpHash preimage: responseCode || commandCode || parameters
+	buf := new(bytes.Buffer)
+
+	// Write responseCode (4 bytes, always 0 for successful responses)
+	if err := binary.Write(buf, binary.BigEndian, uint32(0)); err != nil {
+		return nil, fmt.Errorf("marshalling response code: %w", err)
+	}
+
+	if err := binary.Write(buf, binary.BigEndian, cc); err != nil {
+		return nil, fmt.Errorf("marshalling command code: %w", err)
+	}
+
+	if _, err := buf.Write(params); err != nil {
+		return nil, fmt.Errorf("marshalling parameters: %w", err)
+	}
+
+	return buf.Bytes(), nil
+}
+
+// UnmarshalResponse unmarshals a raw rpHash preimage back into a TPM response.
+// The data should be the output from [MarshalResponse].
+//
+// Example:
+//
+//	rspData, _ := MarshalResponse(commandCode, myRsp)
+//	rsp, _ := UnmarshalResponse[MyResponseType](rspData)
+//
+// Notes:
+//   - the result from this function is expected to be used for purposes such as auditing or inspection.
+//   - encrypted response parameters (via sessions) are not currently supported.
+func UnmarshalResponse[R any](data []byte) (*R, error) {
+	var rsp R
+
+	if data == nil {
+		return nil, fmt.Errorf("data cannot be nil")
+	}
+
+	if len(data) < 8 {
+		return nil, fmt.Errorf("data too short: need at least 8 bytes (responseCode + commandCode), got %d", len(data))
+	}
+
+	buf := bytes.NewBuffer(data)
+
+	var responseCode uint32
+	if err := binary.Read(buf, binary.BigEndian, &responseCode); err != nil {
+		return nil, fmt.Errorf("unmarshalling response code: %w", err)
+	}
+
+	if responseCode != 0 {
+		return nil, fmt.Errorf("invalid response code: expected 0, got 0x%x", responseCode)
+	}
+
+	var cc TPMCC
+	if err := binary.Read(buf, binary.BigEndian, &cc); err != nil {
+		return nil, fmt.Errorf("unmarshalling command code: %w", err)
+	}
+
+	params := buf.Bytes()
+
+	if err := rspParameters(params, nil, &rsp); err != nil {
+		return nil, err
+	}
+	return &rsp, nil
+}

tpm2/marshalling_test.go

@@ -2,7 +2,10 @@ package tpm2
 
 import (
 	"bytes"
+	"reflect"
 	"testing"
+
+	"github.com/google/go-tpm/tpm2/transport/simulator"
 )
 
 func TestMarshal2B(t *testing.T) {
@@ -154,3 +157,48 @@ func TestMarshalT(t *testing.T) {
 		t.Errorf("want %x\ngot %x", pubBytes, pub2Bytes)
 	}
 }
+
+func TestMarshalCommandResponse(t *testing.T) {
+	thetpm, err := simulator.OpenSimulator()
+	if err != nil {
+		t.Fatalf("could not connect to TPM simulator: %v", err)
+	}
+	defer thetpm.Close()
+
+	getCmd := GetCapability{
+		Capability:    TPMCapTPMProperties,
+		Property:      uint32(TPMPTFamilyIndicator),
+		PropertyCount: 1,
+	}
+	capabilityRsp, err := getCmd.Execute(thetpm)
+	if err != nil {
+		t.Fatalf("executing GetCapability: %v", err)
+	}
+
+	cmdParamsBytes, err := MarshalCommand(getCmd)
+	if err != nil {
+		t.Fatalf("MarshalCommand failed: %v", err)
+	}
+
+	unmarshalCmd, err := UnmarshalCommand[GetCapability](cmdParamsBytes)
+	if err != nil {
+		t.Fatalf("UnmarshalCommand failed: %v", err)
+	}
+
+	if !reflect.DeepEqual(getCmd, unmarshalCmd) {
+		t.Errorf("Commands do not match \nwant: %+v\ngot: %+v", getCmd, unmarshalCmd)
+	}
+
+	respParamsBytes, err := MarshalResponse(getCmd, capabilityRsp)
+	if err != nil {
+		t.Fatalf("MarshalResponse failed: %v", err)
+	}
+
+	unmarshalRsp, err := UnmarshalResponse[GetCapabilityResponse](respParamsBytes)
+	if err != nil {
+		t.Fatalf("UnmarshalResponse failed: %v", err)
+	}
+	if !reflect.DeepEqual(capabilityRsp, unmarshalRsp) {
+		t.Errorf("Responses do not match \nwant: %+v\ngot: %+v", capabilityRsp, unmarshalRsp)
+	}
+}