Allow loading of cached keys

client/import.go

@@ -40,7 +40,7 @@ func (k *Key) Import(blob *pb.ImportBlob) ([]byte, error) {
 	}
 	defer tpm2.FlushContext(k.rw, handle)
 
-	unsealSession, err := newPCRSession(k.rw, internal.PCRSelection(blob.Pcrs))
+	unsealSession, err := NewPCRSession(k.rw, internal.PCRSelection(blob.Pcrs))
 	if err != nil {
 		return nil, err
 	}
@@ -76,7 +76,7 @@ func (k *Key) ImportSigningKey(blob *pb.ImportBlob) (key *Key, err error) {
 	if key.pubArea, _, _, err = tpm2.ReadPublic(k.rw, handle); err != nil {
 		return
 	}
-	if key.session, err = newPCRSession(k.rw, internal.PCRSelection(blob.Pcrs)); err != nil {
+	if key.session, err = NewPCRSession(k.rw, internal.PCRSelection(blob.Pcrs)); err != nil {
 		return
 	}
 	return key, key.finish()

client/keys.go

@@ -27,7 +27,7 @@ type Key struct {
 	pubArea tpm2.Public
 	pubKey  crypto.PublicKey
 	name    tpm2.Name
-	session session
+	session Session
 	cert    *x509.Certificate
 }
 
@@ -114,6 +114,19 @@ func GceAttestationKeyECC(rw io.ReadWriter) (*Key, error) {
 	return akEcc, nil
 }
 
+// LoadCachedKey loads a key from cachedHandle.
+// If the key is not found, an error is returned.
+// This function will no overwrite an existing key, unlike NewCachedKey.
+func LoadCachedKey(rw io.ReadWriter, cachedHandle tpmutil.Handle, keySession Session) (k *Key, err error) {
+	cachedPub, _, _, err := tpm2.ReadPublic(rw, cachedHandle)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read public area of cached key: %w", err)
+	}
+
+	k = &Key{rw: rw, handle: cachedHandle, pubArea: cachedPub, session: keySession}
+	return k, k.finish()
+}
+
 // KeyFromNvIndex generates and loads a key under the provided parent
 // (possibly a hierarchy root tpm2.Handle{Owner|Endorsement|Platform|Null})
 // using the template stored at the provided nvdata index.
@@ -182,8 +195,7 @@ func NewKey(rw io.ReadWriter, parent tpmutil.Handle, template tpm2.Public) (k *K
 		return nil, fmt.Errorf("unsupported parent handle: %x", parent)
 	}
 
-	handle, pubArea, _, _, _, _, err :=
-		tpm2.CreatePrimaryEx(rw, parent, tpm2.PCRSelection{}, "", "", template)
+	handle, pubArea, _, _, _, _, err := tpm2.CreatePrimaryEx(rw, parent, tpm2.PCRSelection{}, "", "", template)
 	if err != nil {
 		return nil, err
 	}
@@ -211,11 +223,11 @@ func (k *Key) finish() error {
 	// We determine the right type of session based on the auth policy
 	if k.session == nil {
 		if bytes.Equal(k.pubArea.AuthPolicy, defaultEKAuthPolicy()) {
-			if k.session, err = newEKSession(k.rw); err != nil {
+			if k.session, err = NewEKSession(k.rw); err != nil {
 				return err
 			}
 		} else if len(k.pubArea.AuthPolicy) == 0 {
-			k.session = nullSession{}
+			k.session = NullSession{}
 		} else {
 			return fmt.Errorf("unknown auth policy when creating key")
 		}
@@ -407,7 +419,7 @@ func (k *Key) Unseal(in *pb.SealedBytes, opts UnsealOpts) ([]byte, error) {
 		sel.PCRs = append(sel.PCRs, int(pcr))
 	}
 
-	session, err := newPCRSession(k.rw, sel)
+	session, err := NewPCRSession(k.rw, sel)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create session: %w", err)
 	}

client/session.go

@@ -7,7 +7,7 @@ import (
 	"github.com/google/go-tpm/tpmutil"
 )
 
-type session interface {
+type Session interface {
 	io.Closer
 	Auth() (tpm2.AuthCommand, error)
 }
@@ -31,59 +31,59 @@ func startAuthSession(rw io.ReadWriter) (session tpmutil.Handle, err error) {
 	return
 }
 
-type pcrSession struct {
+type PCRSession struct {
 	rw      io.ReadWriter
 	session tpmutil.Handle
 	sel     tpm2.PCRSelection
 }
 
-func newPCRSession(rw io.ReadWriter, sel tpm2.PCRSelection) (session, error) {
+func NewPCRSession(rw io.ReadWriter, sel tpm2.PCRSelection) (Session, error) {
 	if len(sel.PCRs) == 0 {
-		return nullSession{}, nil
+		return NullSession{}, nil
 	}
 	session, err := startAuthSession(rw)
-	return pcrSession{rw, session, sel}, err
+	return PCRSession{rw, session, sel}, err
 }
 
-func (p pcrSession) Auth() (auth tpm2.AuthCommand, err error) {
+func (p PCRSession) Auth() (auth tpm2.AuthCommand, err error) {
 	if err = tpm2.PolicyPCR(p.rw, p.session, nil, p.sel); err != nil {
 		return
 	}
 	return tpm2.AuthCommand{Session: p.session, Attributes: tpm2.AttrContinueSession}, nil
 }
 
-func (p pcrSession) Close() error {
+func (p PCRSession) Close() error {
 	return tpm2.FlushContext(p.rw, p.session)
 }
 
-type ekSession struct {
+type EKSession struct {
 	rw      io.ReadWriter
 	session tpmutil.Handle
 }
 
-func newEKSession(rw io.ReadWriter) (session, error) {
+func NewEKSession(rw io.ReadWriter) (Session, error) {
 	session, err := startAuthSession(rw)
-	return ekSession{rw, session}, err
+	return EKSession{rw, session}, err
 }
 
-func (e ekSession) Auth() (auth tpm2.AuthCommand, err error) {
+func (e EKSession) Auth() (auth tpm2.AuthCommand, err error) {
 	nullAuth := tpm2.AuthCommand{Session: tpm2.HandlePasswordSession, Attributes: tpm2.AttrContinueSession}
 	if _, _, err = tpm2.PolicySecret(e.rw, tpm2.HandleEndorsement, nullAuth, e.session, nil, nil, nil, 0); err != nil {
 		return
 	}
 	return tpm2.AuthCommand{Session: e.session, Attributes: tpm2.AttrContinueSession}, nil
 }
 
-func (e ekSession) Close() error {
+func (e EKSession) Close() error {
 	return tpm2.FlushContext(e.rw, e.session)
 }
 
-type nullSession struct{}
+type NullSession struct{}
 
-func (n nullSession) Auth() (auth tpm2.AuthCommand, err error) {
+func (n NullSession) Auth() (auth tpm2.AuthCommand, err error) {
 	return tpm2.AuthCommand{Session: tpm2.HandlePasswordSession, Attributes: tpm2.AttrContinueSession}, nil
 }
 
-func (n nullSession) Close() error {
+func (n NullSession) Close() error {
 	return nil
 }