From 8cae640078ff661223d3d02c26748e52f790b660 Mon Sep 17 00:00:00 2001
From: Alex Wu <wuale@google.com>
Date: Mon, 8 Mar 2021 16:28:15 -0800
Subject: [PATCH] Add support for encoding Signatures

Some functions in tpm2 expect an encoded TPMT_SIGNATURE.
Here, we add an Encode method on the Signature type
to support these functions.
---
 tpm2/structures.go     | 21 ++++++++++++++
 tpm2/test/tpm2_test.go | 65 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 86 insertions(+)

diff --git a/tpm2/structures.go b/tpm2/structures.go
index da81dc77..6df9f7f0 100644
--- a/tpm2/structures.go
+++ b/tpm2/structures.go
@@ -573,6 +573,27 @@ type Signature struct {
 	ECC *SignatureECC
 }
 
+// Encode serializes a Signature structure in TPM wire format.
+func (s Signature) Encode() ([]byte, error) {
+	head, err := tpmutil.Pack(s.Alg)
+	if err != nil {
+		return nil, fmt.Errorf("encoding Alg: %v", err)
+	}
+	var signature []byte
+	switch s.Alg {
+	case AlgRSASSA, AlgRSAPSS:
+		if signature, err = tpmutil.Pack(s.RSA); err != nil {
+			return nil, fmt.Errorf("encoding RSA: %v", err)
+		}
+	case AlgECDSA:
+		signature, err = tpmutil.Pack(s.ECC.HashAlg, tpmutil.U16Bytes(s.ECC.R.Bytes()), tpmutil.U16Bytes(s.ECC.S.Bytes()))
+		if err != nil {
+			return nil, fmt.Errorf("encoding ECC: %v", err)
+		}
+	}
+	return concat(head, signature)
+}
+
 // DecodeSignature decodes a serialized TPMT_SIGNATURE structure.
 func DecodeSignature(in *bytes.Buffer) (*Signature, error) {
 	var sig Signature
diff --git a/tpm2/test/tpm2_test.go b/tpm2/test/tpm2_test.go
index 3a845a52..cad9e8ac 100644
--- a/tpm2/test/tpm2_test.go
+++ b/tpm2/test/tpm2_test.go
@@ -23,10 +23,12 @@ import (
 	"crypto/rsa"
 	"crypto/sha1"
 	"crypto/sha256"
+	"encoding/binary"
 	"flag"
 	"fmt"
 	"hash"
 	"io"
+	"math/big"
 	"reflect"
 	"strings"
 	"testing"
@@ -1179,6 +1181,69 @@ func TestEncodeDecodePublicDefaultRSAExponent(t *testing.T) {
 	}
 }
 
+func TestEncodeDecodeSignature(t *testing.T) {
+	randRSASig := func() []byte {
+		// Key size 2048 bits
+		var size uint16 = 256
+		sizeU16 := make([]byte, 2)
+		binary.BigEndian.PutUint16(sizeU16, size)
+		key := make([]byte, size)
+		rand.Read(key)
+		return append(sizeU16, key...)
+	}
+
+	run := func(t *testing.T, s Signature) {
+		e, err := s.Encode()
+		if err != nil {
+			t.Fatalf("Signature{%+v}.Encode() returned error: %v", s, err)
+		}
+		d, err := DecodeSignature(bytes.NewBuffer(e))
+		if err != nil {
+			t.Fatalf("DecodeSignature{%v} returned error: %v", e, err)
+		}
+		if !reflect.DeepEqual(s, *d) {
+			t.Errorf("got decoded value:\n%v\nwant:\n%v", d, s)
+		}
+	}
+	t.Run("RSASSA", func(t *testing.T) {
+		run(t, Signature{
+			Alg: AlgRSASSA,
+			RSA: &SignatureRSA{
+				HashAlg:   AlgSHA256,
+				Signature: randRSASig(),
+			},
+		})
+	})
+	t.Run("RSAPSS", func(t *testing.T) {
+		run(t, Signature{
+			Alg: AlgRSAPSS,
+			RSA: &SignatureRSA{
+				HashAlg:   AlgSHA256,
+				Signature: randRSASig(),
+			},
+		})
+	})
+	t.Run("ECDSA", func(t *testing.T) {
+		// Key size 256 bits
+		size := 32
+		randBytes := make([]byte, size)
+		rand.Read(randBytes)
+		r := big.NewInt(0).SetBytes(randBytes)
+
+		rand.Read(randBytes)
+		s := big.NewInt(0).SetBytes(randBytes)
+
+		run(t, Signature{
+			Alg: AlgECDSA,
+			ECC: &SignatureECC{
+				HashAlg: AlgSHA256,
+				R:       r,
+				S:       s,
+			},
+		})
+	})
+}
+
 func TestCreateKeyWithSensitive(t *testing.T) {
 	rw := openTPM(t)
 	defer rw.Close()