Abstract public key handling

client_test.go

@@ -204,7 +204,7 @@ func TestDoUpdate(t *testing.T) {
 	require.NoError(t, err)
 
 	// convert privkey to private key
-	pubkey, err := keys.MarshalEd25519HostPublicKey(creds.PrivateKey.Unwrap().(ed25519.PrivateKey).Public().(ed25519.PublicKey))
+	pubkey, err := keys.MarshalEd25519HostPublicKey(creds.PrivateKey.Public().Unwrap().(ed25519.PublicKey))
 	require.NoError(t, err)
 
 	// make sure all credential values were set
@@ -395,7 +395,7 @@ func TestDoUpdate_P256(t *testing.T) {
 	require.NoError(t, err)
 
 	// convert privkey to private key
-	pubkey, err := keys.MarshalP256HostPublicKey(creds.PrivateKey.Unwrap().(*ecdsa.PrivateKey).Public().(*ecdsa.PublicKey))
+	pubkey, err := keys.MarshalP256HostPublicKey(creds.PrivateKey.Public().Unwrap().(*ecdsa.PublicKey))
 	require.NoError(t, err)
 
 	// make sure all credential values were set

keys/crypto.go

@@ -13,6 +13,27 @@ import (
 	"golang.org/x/crypto/curve25519"
 )
 
+// PublicKey is a wrapper around public keys.
+type PublicKey interface {
+	Unwrap() interface{}
+}
+
+type Ed25519PublicKey struct {
+	ed25519.PublicKey
+}
+
+func (k Ed25519PublicKey) Unwrap() interface{} {
+	return k.PublicKey
+}
+
+type P256PublicKey struct {
+	*ecdsa.PublicKey
+}
+
+func (k P256PublicKey) Unwrap() interface{} {
+	return k.PublicKey
+}
+
 // PrivateKey is an interface used to generically sign messages regardless of
 // the network curve (P256/25519.)
 type PrivateKey interface {
@@ -24,6 +45,9 @@ type PrivateKey interface {
 
 	// MarshalPEM returns the private key in PEM format.
 	MarshalPEM() ([]byte, error)
+
+	// Public returns the public key associated with the private key.
+	Public() PublicKey
 }
 
 func NewPrivateKey(k any) (PrivateKey, error) {
@@ -37,6 +61,8 @@ func NewPrivateKey(k any) (PrivateKey, error) {
 	}
 }
 
+// Ed25519PrivateKey is a wrapper around an Ed25519 private key that implements
+// the PrivateKey interface.
 type Ed25519PrivateKey struct {
 	ed25519.PrivateKey
 }
@@ -53,6 +79,12 @@ func (k Ed25519PrivateKey) MarshalPEM() ([]byte, error) {
 	return MarshalEd25519HostPrivateKey(k.PrivateKey)
 }
 
+func (k Ed25519PrivateKey) Public() PublicKey {
+	return Ed25519PublicKey{k.PrivateKey.Public().(ed25519.PublicKey)}
+}
+
+// P256PrivateKey is a wrapper around an ECDSA private key that implements the
+// PrivateKey interface.
 type P256PrivateKey struct {
 	*ecdsa.PrivateKey
 }
@@ -70,6 +102,10 @@ func (k P256PrivateKey) MarshalPEM() ([]byte, error) {
 	return MarshalP256HostPrivateKey(k.PrivateKey)
 }
 
+func (k P256PrivateKey) Public() PublicKey {
+	return P256PublicKey{k.PrivateKey.Public().(*ecdsa.PublicKey)}
+}
+
 // Keys contains a set of P256 and X25519/Ed25519 keys. Only one set is used,
 // depending on the network the host is enrolled in. At the time of enrollment
 // clients do not know which curve the network uses, so both keys must be