wip: include PR's feedbacks

Author: loicsikidi

tpm2/marshalling.go

@@ -2,6 +2,7 @@ package tpm2
 
 import (
 	"bytes"
+	"encoding/binary"
 	"fmt"
 	"reflect"
 )
@@ -22,6 +23,45 @@ type marshallableWithHint interface {
 	get(hint int64) (reflect.Value, error)
 }
 
+// CommandPreimage represents a structured preimage of cpHash for a TPM command.
+// This structure is marshaled to bytes using [Marshal] for storage/transmission
+// and can be converted to the raw cpHash preimage format for hashing.
+//
+// Format when marshaled:
+//   - CommandCode: 4 bytes (TPMCC)
+//   - Names: sized list of TPM2BName
+//   - Parameters: sized buffer
+//
+// See definition in Part 1: Architecture, section 16.7.
+type CommandPreimage struct {
+	marshalByReflection
+	// CommandCode is the TPM command code
+	CommandCode TPMCC
+	// Names are the names of the handles referenced by the command
+	Names []TPM2BName `gotpm:"list"`
+	// Parameters are the marshaled command parameters
+	Parameters TPM2BData
+}
+
+// ToCPHashPreimage converts the CommandPreimage to the raw buffer format
+// used to compute cpHash according to TPM 2.0 spec.
+func (cp *CommandPreimage) ToCPHashPreimage() []byte {
+	var buf bytes.Buffer
+
+	// Write command code (4 bytes, big endian)
+	binary.Write(&buf, binary.BigEndian, cp.CommandCode)
+
+	// Write names (raw buffers without size prefix)
+	for _, name := range cp.Names {
+		buf.Write(name.Buffer)
+	}
+
+	// Write parameters
+	buf.Write(cp.Parameters.Buffer)
+
+	return buf.Bytes()
+}
+
 // Unmarshallable represents any TPM type that can be marshalled or unmarshalled.
 type Unmarshallable interface {
 	Marshallable
@@ -127,84 +167,118 @@ func (b *boxed[T]) unmarshal(buf *bytes.Buffer) error {
 	return unmarshal(buf, reflect.ValueOf(b.Contents))
 }
 
-// MarshalCommandResponse marshals both command and response.
-func MarshalCommandResponse[C Command[R, *R], R any](cmd C, rsp *R) (cmdData []byte, rspData []byte, err error) {
-	cmdData, err = MarshalCommand(cmd)
+// toCommandPreimage convert a Command to a CommandPreimage structure from a command.
+func toCommandPreimage[C Command[R, *R], R any](cmd C) (*CommandPreimage, error) {
+	cc := cmd.Command()
+
+	names, err := cmdNames(cmd)
 	if err != nil {
-		return nil, nil, fmt.Errorf("marshalling command: %w", err)
+		return nil, err
 	}
-	rspData, err = MarshalResponse(rsp)
+
+	params, err := cmdParameters(cmd, nil)
 	if err != nil {
-		return nil, nil, fmt.Errorf("marshalling response: %w", err)
+		return nil, err
 	}
-	return cmdData, rspData, nil
+
+	return &CommandPreimage{
+		CommandCode: cc,
+		Names:       names,
+		Parameters: TPM2BData{
+			Buffer: params,
+		},
+	}, nil
 }
 
-// UnmarshalCommandResponse unmarshals both command and response.
-func UnmarshalCommandResponse[C Command[R, *R], R any](cmdData []byte, rspData []byte) (cmd C, rsp *R, err error) {
-	cmd, err = UnmarshalCommand[C, R](cmdData)
-	if err != nil {
-		return cmd, rsp, fmt.Errorf("unmarshalling command: %w", err)
-	}
-	rsp, err = UnmarshalResponse[R](rspData)
+// MarshalCommand marshals a TPM command into a serialized CommandPreimage.
+// The returned bytes contain a marshaled CommandPreimage structure that includes:
+//   - CommandCode (4 bytes)
+//   - Names (sized list)
+//   - Parameters (sized buffer)
+//
+// This can be stored, transmitted, or later unmarshaled.
+//
+// To compute cpHash use [CommandPreimage.ToCPHashPreimage].
+func MarshalCommand[C Command[R, *R], R any](cmd C) ([]byte, error) {
+	preimage, err := toCommandPreimage(cmd)
 	if err != nil {
-		return cmd, rsp, fmt.Errorf("unmarshalling response: %w", err)
+		return nil, err
 	}
-	return cmd, rsp, nil
+	return Marshal(preimage), nil
 }
 
-// MarshalCommand marshals a TPM command.
-func MarshalCommand[C Command[R, *R], R any](cmd C) ([]byte, error) {
-	var buf bytes.Buffer
-	params := taggedMembers(reflect.ValueOf(cmd), "handle", true)
-	for i := range len(params) {
-		if err := marshalParameter(&buf, cmd, i); err != nil {
-			return nil, fmt.Errorf("marshalling command's parameter: %w", err)
-		}
+// unmarshalCommandPreimage unmarshals serialized data into CommandPreimage components.
+// Returns the command code, names, and parameters.
+func unmarshalCommandPreimage(data []byte) (TPMCC, []TPM2BName, []byte, error) {
+	if data == nil {
+		return 0, nil, nil, fmt.Errorf("data cannot be nil")
 	}
-	return buf.Bytes(), nil
+
+	preimage, err := Unmarshal[CommandPreimage](data)
+	if err != nil {
+		return 0, nil, nil, fmt.Errorf("unmarshalling CommandPreimage: %w", err)
+	}
+
+	return preimage.CommandCode, preimage.Names, preimage.Parameters.Buffer, nil
 }
 
-// UnmarshalCommand unmarshals a TPM command.
+// UnmarshalCommand unmarshals a serialized [CommandPreimage] back into a TPM command.
+// The data should be the output from [MarshalCommand].
+//
+// Note: 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.
 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")
+
+	cc, names, params, err := unmarshalCommandPreimage(data)
+	if err != nil {
+		return cmd, err
 	}
-	buf := bytes.NewBuffer(data)
-	params := taggedMembers(reflect.ValueOf(cmd), "handle", true)
-	for i := range len(params) {
-		if err := unmarshalParameter(buf, &cmd, i); err != nil {
-			return cmd, fmt.Errorf("unmarshalling command's parameter: %w", err)
+
+	if cc != cmd.Command() {
+		return cmd, fmt.Errorf("command code mismatch: expected %v, got %v", cmd.Command(), cc)
+	}
+
+	{
+		n, err := cmdNames(cmd)
+		if err != nil {
+			return cmd, err
 		}
+
+		expectedNames := len(names)
+		if len(n) != expectedNames {
+			return cmd, fmt.Errorf("name count mismatch: command expects %d names, got %d", expectedNames, len(n))
+		}
+	}
+
+	// Populate the command's handle fields from the names
+	if err := populateHandlesFromNames(&cmd, names); err != nil {
+		return cmd, fmt.Errorf("populating handles: %w", err)
+	}
+
+	// Now unmarshal the parameters using the helper
+	buf := bytes.NewBuffer(params)
+	if err := unmarshalCmdParameters(buf, &cmd, nil); err != nil {
+		return cmd, err
 	}
 	return cmd, nil
 }
 
 // MarshalResponse marshals a TPM response.
 func MarshalResponse[R any](rsp *R) ([]byte, error) {
-	var buf bytes.Buffer
-	parameters := taggedMembers(reflect.ValueOf(rsp).Elem(), "handle", true)
-	for i, parameter := range parameters {
-		if err := marshal(&buf, parameter); err != nil {
-			return nil, fmt.Errorf("marshalling response parameter %d: %w", i, err)
-		}
-	}
-	return buf.Bytes(), nil
+	return marshalRspParameters(rsp, nil)
 }
 
 // UnmarshalResponse unmarshals a TPM response.
+//
+// Note: the result from this function is expected to be used for purposes such as auditing or inspection.
 func UnmarshalResponse[R any](data []byte) (*R, error) {
 	var rsp R
 	if data == nil {
 		return nil, fmt.Errorf("data cannot be nil")
 	}
-	buf := bytes.NewBuffer(data)
-	parameters := taggedMembers(reflect.ValueOf(&rsp).Elem(), "handle", true)
-	for i, parameter := range parameters {
-		if err := unmarshal(buf, parameter); err != nil {
-			return nil, fmt.Errorf("unmarshalling response parameter %d: %w", i, err)
-		}
+	if err := rspParameters(data, nil, &rsp); err != nil {
+		return nil, err
 	}
 	return &rsp, nil
 }