Use generics to simplify the TPMDirect interface

.cirrus.yml

@@ -27,6 +27,7 @@ lint_task:
       --exclude-use-default=false
       --exclude stutters
       --exclude underscores
+      --exclude unexported-return
       --max-same-issues=0
       --max-issues-per-linter=0
       ./tpmutil/...

go.mod

@@ -3,9 +3,9 @@ module github.com/google/go-tpm
 go 1.18
 
 require (
-	github.com/google/go-cmp v0.5.0
-	github.com/google/go-tpm-tools v0.2.0
-	golang.org/x/sys v0.0.0-20210629170331-7dc0b73dc9fb
+	github.com/google/go-cmp v0.5.7
+	github.com/google/go-tpm-tools v0.3.10
+	golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f
 )
 
-require golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 // indirect
+require golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect

tpm2/audit.go

@@ -78,7 +78,7 @@ func auditRPHash(h TPMIAlgHash, r Response) ([]byte, error) {
 	parameters := taggedMembers(reflect.ValueOf(r).Elem(), "handle", true)
 	for i, parameter := range parameters {
 		if err := marshal(&parms, parameter); err != nil {
-			return nil, fmt.Errorf("marshalling parameter %v: %w", i, err)
+			return nil, fmt.Errorf("marshalling parameter %v: %w", i+1, err)
 		}
 	}
 	return rpHash(h, TPMRCSuccess, cc, parms.Bytes())

tpm2/marshalling.go

@@ -0,0 +1,191 @@
+package tpm2
+
+import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
+	"io"
+	"reflect"
+)
+
+// Marshallable represents any TPM type that can be marshalled.
+type Marshallable interface {
+	// marshal will serialize the given value, appending onto the given buffer.
+	// Returns an error if the value is not marshallable.
+	marshal(buf *bytes.Buffer)
+}
+
+// Unmarshallable represents any TPM type that can be marshalled or unmarshalled.
+type Unmarshallable interface {
+	Marshallable
+	// marshal will deserialize the given value from the given buffer.
+	// Returns an error if there was an unmarshalling error or if there was not
+	// enough data in the buffer.
+	unmarshal(buf *bytes.Buffer) error
+}
+
+// unmarshallableWithHint represents any TPM type that can be marshalled or unmarshalled,
+// but that requires a selector ("hint") value when unmarshalling. Most TPMU_ are
+// an example of this.
+type unmarshallableWithHint interface {
+	// create will instantiate and return the corresponding union member.
+	create(hint int64) (reflect.Value, error)
+	// get will return the corresponding union member by copy. If the union is
+	// uninitialized, it will initialize a new zero-valued one.
+	get(hint int64) (reflect.Value, error)
+}
+
+// Marshal will serialize the given values, returning them as a byte slice.
+func Marshal(v Marshallable) []byte {
+	var buf bytes.Buffer
+	if err := marshal(&buf, reflect.ValueOf(v)); err != nil {
+		panic(fmt.Sprintf("unexpected error marshalling %v: %v", reflect.TypeOf(v).Name(), err))
+	}
+	return buf.Bytes()
+}
+
+// Unmarshal unmarshals the given type from the byte array.
+// Returns an error if the buffer does not contain enough data to satisfy the
+// types, or if the types are not unmarshallable.
+func Unmarshal[T any, P interface {
+	// *T must satisfy Marshallable
+	*T
+	Unmarshallable
+}](data []byte) (*T, error) {
+	buf := bytes.NewBuffer(data)
+	var t T
+	value := reflect.New(reflect.TypeOf(t))
+	if err := unmarshal(buf, value.Elem()); err != nil {
+		return nil, err
+	}
+	return value.Interface().(*T), nil
+}
+
+// marshallableByReflection is a placeholder interface, to hint to the unmarshalling
+// library that it is supposed to use reflection.
+type marshallableByReflection interface {
+	reflectionSafe()
+}
+
+// marshalByReflection is embedded into any type that can be marshalled by reflection,
+// needing no custom logic.
+type marshalByReflection struct{}
+
+func (marshalByReflection) reflectionSafe() {}
+
+// These placeholders are required because a type constraint cannot union another interface
+// that contains methods.
+// Otherwise, marshalByReflection would not implement Unmarshallable, and the Marshal/Unmarshal
+// functions would accept interface{ Marshallable | marshallableByReflection } instead.
+
+// Placeholder: because this type implements the defaultMarshallable interface,
+// the reflection library knows not to call this.
+func (*marshalByReflection) marshal(_ *bytes.Buffer) {
+	panic("not implemented")
+}
+
+// Placeholder: because this type implements the defaultMarshallable interface,
+// the reflection library knows not to call this.
+func (*marshalByReflection) unmarshal(_ *bytes.Buffer) error {
+	panic("not implemented")
+}
+
+// tpm2b is a helper type for a field that contains either a structure or by byte-array.
+// When serialized, if contents is non-nil, the value of contents is used.
+// Because it can be instantiated by value or by byte-array, Marshal must be used in order to get
+// the flattened data.
+//
+//	Else, the value of Buffer is used.
+type tpm2b[T any] struct {
+	contents *T
+	buffer   []byte
+}
+
+type bytesOr[T any] interface{ *T | []byte }
+
+// tpm2bHelper is a helper function that can convert either a structure or a byte buffer into
+// the proper TPM2B_ sub-type.
+func tpm2bHelper[T any, C bytesOr[T]](contents C) tpm2b[T] {
+	if typed, ok := any(contents).(*T); ok {
+		return tpm2b[T]{contents: typed}
+	}
+	return tpm2b[T]{buffer: any(contents).([]byte)}
+}
+
+// marshal implements the Marshallable interface.
+func (value *tpm2b[T]) marshal(buf *bytes.Buffer) {
+	if value.contents != nil {
+		var temp bytes.Buffer
+		marshal(&temp, reflect.ValueOf(value.contents))
+		binary.Write(buf, binary.BigEndian, uint16(temp.Len()))
+		io.Copy(buf, &temp)
+	} else {
+		binary.Write(buf, binary.BigEndian, uint16(len(value.buffer)))
+		buf.Write(value.buffer)
+	}
+}
+
+// unmarshal implements the Marshallable interface.
+func (value *tpm2b[T]) unmarshal(buf *bytes.Buffer) error {
+	var size uint16
+	binary.Read(buf, binary.BigEndian, &size)
+	value.buffer = make([]byte, size)
+	n, err := buf.Read(value.buffer)
+	if err != nil {
+		return err
+	}
+	if n != int(size) {
+		return fmt.Errorf("ran out of data attempting to read %v bytes from the buffer, which only had %v", size, n)
+	}
+	rdr := bytes.NewBuffer(value.buffer)
+	value.contents = new(T)
+	return unmarshal(rdr, reflect.ValueOf(value.contents))
+}
+
+// Contents returns the structured contents of the tpm2b.
+func (value *tpm2b[T]) Contents() (*T, error) {
+	if value.contents != nil {
+		return value.contents, nil
+	}
+	if value.buffer == nil {
+		return nil, fmt.Errorf("TPMB had no contents or buffer")
+	}
+	var result T
+	if err := unmarshal(bytes.NewBuffer(value.buffer), reflect.ValueOf(&result).Elem()); err != nil {
+		return nil, err
+	}
+	return &result, nil
+}
+
+// boxed is a helper type for corner cases such as unions, where all members must be structs.
+type boxed[T any] struct {
+	Contents *T
+}
+
+// box will put a value into a box.
+func box[T any](contents *T) boxed[T] {
+	return boxed[T]{
+		Contents: contents,
+	}
+}
+
+// unbox will take a value out of a box.
+func (b *boxed[T]) unbox() *T {
+	return b.Contents
+}
+
+// marshal implements the Marshallable interface.
+func (b *boxed[T]) marshal(buf *bytes.Buffer) {
+	if b.Contents == nil {
+		var contents T
+		marshal(buf, reflect.ValueOf(&contents))
+	} else {
+		marshal(buf, reflect.ValueOf(b.Contents))
+	}
+}
+
+// unmarshal implements the Unmarshallable interface.
+func (b *boxed[T]) unmarshal(buf *bytes.Buffer) error {
+	b.Contents = new(T)
+	return unmarshal(buf, reflect.ValueOf(b.Contents))
+}

tpm2/marshalling_test.go

@@ -0,0 +1,156 @@
+package tpm2
+
+import (
+	"bytes"
+	"testing"
+)
+
+func TestMarshal2B(t *testing.T) {
+	// Define some TPMT_Public
+	pub := TPMTPublic{
+		Type:    TPMAlgKeyedHash,
+		NameAlg: TPMAlgSHA256,
+		ObjectAttributes: TPMAObject{
+			FixedTPM:     true,
+			FixedParent:  true,
+			UserWithAuth: true,
+			NoDA:         true,
+		},
+	}
+
+	// Get the wire-format version
+	pubBytes := Marshal(&pub)
+
+	// Create two versions of the same 2B:
+	// one instantiated by the actual TPMTPublic
+	// one instantiated by the contents
+	var boxed1 tpm2bPublic
+	var boxed2 tpm2bPublic
+	boxed1 = TPM2BPublic(&pub)
+	boxed2 = TPM2BPublic(pubBytes)
+
+	boxed1Bytes := Marshal(&boxed1)
+	boxed2Bytes := Marshal(&boxed2)
+
+	if !bytes.Equal(boxed1Bytes, boxed2Bytes) {
+		t.Errorf("got %x want %x", boxed2Bytes, boxed1Bytes)
+	}
+
+	z, err := Unmarshal[tpm2bPublic](boxed1Bytes)
+	if err != nil {
+		t.Fatalf("could not unmarshal TPM2BPublic: %v", err)
+	}
+	t.Logf("%v", z)
+
+	boxed3Bytes := Marshal(z)
+	if !bytes.Equal(boxed1Bytes, boxed3Bytes) {
+		t.Errorf("got %x want %x", boxed3Bytes, boxed1Bytes)
+	}
+
+	// Make a nonsense 2B_Public, demonstrating that the library doesn't have to understand the serialization
+	boxed1 = TPM2BPublic([]byte{0xff})
+}
+
+func unwrap[T any](f func() (*T, error)) *T {
+	t, err := f()
+	if err != nil {
+		panic(err.Error())
+	}
+	return t
+}
+
+func TestMarshalT(t *testing.T) {
+	// Define some TPMT_Public
+	pub := TPMTPublic{
+		Type:    TPMAlgECC,
+		NameAlg: TPMAlgSHA256,
+		ObjectAttributes: TPMAObject{
+			SignEncrypt: true,
+		},
+		Parameters: TPMUPublicParms(
+			TPMAlgECC,
+			&TPMSECCParms{
+				CurveID: TPMECCNistP256,
+			},
+		),
+		Unique: TPMUPublicID(
+			// This happens to be a P256 EKpub from the simulator
+			TPMAlgECC,
+			&TPMSECCPoint{
+				X: TPM2BECCParameter{},
+				Y: TPM2BECCParameter{},
+			},
+		),
+	}
+
+	// Marshal each component of the parameters
+	symBytes := Marshal(&unwrap(pub.Parameters.ECCDetail).Symmetric)
+	t.Logf("Symmetric: %x\n", symBytes)
+	sym, err := Unmarshal[TPMTSymDefObject](symBytes)
+	if err != nil {
+		t.Fatalf("could not unmarshal TPMTSymDefObject: %v", err)
+	}
+	symBytes2 := Marshal(sym)
+	if !bytes.Equal(symBytes, symBytes2) {
+		t.Errorf("want %x\ngot %x", symBytes, symBytes2)
+	}
+	schemeBytes := Marshal(&unwrap(pub.Parameters.ECCDetail).Scheme)
+	t.Logf("Scheme: %x\n", symBytes)
+	scheme, err := Unmarshal[TPMTECCScheme](schemeBytes)
+	if err != nil {
+		t.Fatalf("could not unmarshal TPMTECCScheme: %v", err)
+	}
+	schemeBytes2 := Marshal(scheme)
+	if !bytes.Equal(schemeBytes, schemeBytes2) {
+		t.Errorf("want %x\ngot %x", schemeBytes, schemeBytes2)
+	}
+	kdfBytes := Marshal(&unwrap(pub.Parameters.ECCDetail).KDF)
+	t.Logf("KDF: %x\n", kdfBytes)
+	kdf, err := Unmarshal[TPMTKDFScheme](kdfBytes)
+	if err != nil {
+		t.Fatalf("could not unmarshal TPMTKDFScheme: %v", err)
+	}
+	kdfBytes2 := Marshal(kdf)
+	if !bytes.Equal(kdfBytes, kdfBytes2) {
+		t.Errorf("want %x\ngot %x", kdfBytes, kdfBytes2)
+	}
+
+	// Marshal the parameters
+	parmsBytes := Marshal(unwrap(pub.Parameters.ECCDetail))
+	t.Logf("Parms: %x\n", parmsBytes)
+	parms, err := Unmarshal[TPMSECCParms](parmsBytes)
+	if err != nil {
+		t.Fatalf("could not unmarshal TPMSECCParms: %v", err)
+	}
+	parmsBytes2 := Marshal(parms)
+	if !bytes.Equal(parmsBytes, parmsBytes2) {
+		t.Errorf("want %x\ngot %x", parmsBytes, parmsBytes2)
+	}
+
+	// Marshal the unique area
+	uniqueBytes := Marshal(unwrap(pub.Unique.ECC))
+	t.Logf("Unique: %x\n", uniqueBytes)
+	unique, err := Unmarshal[TPMSECCPoint](uniqueBytes)
+	if err != nil {
+		t.Fatalf("could not unmarshal TPMSECCPoint: %v", err)
+	}
+	uniqueBytes2 := Marshal(unique)
+	if !bytes.Equal(uniqueBytes, uniqueBytes2) {
+		t.Errorf("want %x\ngot %x", uniqueBytes, uniqueBytes2)
+	}
+
+	// Get the wire-format version of the whole thing
+	pubBytes := Marshal(&pub)
+
+	pub2, err := Unmarshal[TPMTPublic](pubBytes)
+	if err != nil {
+		t.Fatalf("could not unmarshal TPMTPublic: %v", err)
+	}
+
+	// Some default fields might have been populated in the round-trip. Get the wire-format again and compare.
+	pub2Bytes := Marshal(pub2)
+
+	if !bytes.Equal(pubBytes, pub2Bytes) {
+		t.Errorf("want %x\ngot %x", pubBytes, pub2Bytes)
+	}
+}