Return wrapped type to avoid callers having to do it

DefinedNet · Dec 5, 2024 · 90c503c · 90c503c
1 parent da86003
commit 90c503c
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 35 deletions.
diff --git a/client.go b/client.go
@@ -169,16 +169,10 @@ func (c *Client) Enroll(ctx context.Context, logger logrus.FieldLogger, code str
 	switch r.Data.Network.Curve {
 	case message.NetworkCurve25519:
 		privkeyPEM = newKeys.NebulaX25519PrivateKeyPEM
-		privkey, err = keys.NewPrivateKey(newKeys.HostEd25519PrivateKey)
-		if err != nil {
-			return nil, nil, nil, nil, &APIError{e: fmt.Errorf("failed to create new private key: %s", err), ReqID: reqID}
-		}
+		privkey = newKeys.HostEd25519PrivateKey
 	case message.NetworkCurveP256:
 		privkeyPEM = newKeys.NebulaP256PrivateKeyPEM
-		privkey, err = keys.NewPrivateKey(newKeys.HostP256PrivateKey)
-		if err != nil {
-			return nil, nil, nil, nil, &APIError{e: fmt.Errorf("failed to create new private key: %s", err), ReqID: reqID}
-		}
+		privkey = newKeys.HostP256PrivateKey
 	default:
 		return nil, nil, nil, nil, &APIError{e: fmt.Errorf("unsupported curve type: %s", r.Data.Network.Curve), ReqID: reqID}
 	}
@@ -255,18 +249,12 @@ func (c *Client) DoUpdate(ctx context.Context, creds keys.Credentials) ([]byte,
 	// Set the correct keypair based on the current private key type
 	switch creds.PrivateKey.Unwrap().(type) {
 	case ed25519.PrivateKey:
-		hostPrivkey, err = keys.NewPrivateKey(newKeys.HostEd25519PrivateKey)
-		if err != nil {
-			return nil, nil, nil, fmt.Errorf("failed to create new private key: %s", err)
-		}
+		hostPrivkey = newKeys.HostEd25519PrivateKey
 		nebulaPrivkeyPEM = newKeys.NebulaX25519PrivateKeyPEM
 		msg.HostPubkeyEd25519 = newKeys.HostEd25519PublicKeyPEM
 		msg.NebulaPubkeyX25519 = newKeys.NebulaX25519PublicKeyPEM
 	case *ecdsa.PrivateKey:
-		hostPrivkey, err = keys.NewPrivateKey(newKeys.HostP256PrivateKey)
-		if err != nil {
-			return nil, nil, nil, fmt.Errorf("failed to create new private key: %s", err)
-		}
+		hostPrivkey = newKeys.HostP256PrivateKey
 		nebulaPrivkeyPEM = newKeys.NebulaP256PrivateKeyPEM
 		msg.HostPubkeyP256 = newKeys.HostP256PublicKeyPEM
 		msg.NebulaPubkeyP256 = newKeys.NebulaP256PublicKeyPEM

diff --git a/client_test.go b/client_test.go
@@ -225,12 +225,10 @@ func TestDoUpdate(t *testing.T) {
 	// Create a new, invalid requesting authentication key
 	nk, err := keys.New()
 	require.NoError(t, err)
-	pk, err := keys.NewPrivateKey(nk.HostEd25519PrivateKey)
-	require.NoError(t, err)
 
 	invalidCreds := keys.Credentials{
 		HostID:      creds.HostID,
-		PrivateKey:  pk,
+		PrivateKey:  nk.HostEd25519PrivateKey,
 		Counter:     creds.Counter,
 		TrustedKeys: creds.TrustedKeys,
 	}
@@ -261,11 +259,14 @@ func TestDoUpdate(t *testing.T) {
 		err = ts.SetEdPubkey(pubkey)
 		require.NoError(t, err)
 
+		sig, err := nk.HostEd25519PrivateKey.Sign(rawRes)
+		require.NoError(t, err)
+
 		return jsonMarshal(message.SignedResponseWrapper{
 			Data: message.SignedResponse{
 				Version:   1,
 				Message:   rawRes,
-				Signature: ed25519.Sign(nk.HostEd25519PrivateKey, rawRes),
+				Signature: sig,
 			},
 		})
 	})
@@ -415,11 +416,9 @@ func TestDoUpdate_P256(t *testing.T) {
 	// Create a new, invalid requesting authentication key
 	nk, err := keys.New()
 	require.NoError(t, err)
-	pk, err := keys.NewPrivateKey(nk.HostP256PrivateKey)
-	require.NoError(t, err)
 	invalidCreds := keys.Credentials{
 		HostID:      creds.HostID,
-		PrivateKey:  pk,
+		PrivateKey:  nk.HostP256PrivateKey,
 		Counter:     creds.Counter,
 		TrustedKeys: creds.TrustedKeys,
 	}
@@ -450,8 +449,7 @@ func TestDoUpdate_P256(t *testing.T) {
 		err = ts.SetP256Pubkey(pubkey)
 		require.NoError(t, err)
 
-		hashed := sha256.Sum256(rawRes)
-		sig, err := ecdsa.SignASN1(rand.Reader, nk.HostP256PrivateKey, hashed[:])
+		sig, err := nk.HostP256PrivateKey.Sign(rawRes)
 		if err != nil {
 			return jsonMarshal(message.EnrollResponse{
 				Errors: message.APIErrors{{

diff --git a/keys/crypto.go b/keys/crypto.go
@@ -81,16 +81,16 @@ func (k P256PrivateKey) MarshalPEM() ([]byte, error) {
 // we will need it to sign requests.
 type Keys struct {
 	// 25519 Curve
-	NebulaX25519PublicKeyPEM  []byte             // ECDH (Nebula)
-	NebulaX25519PrivateKeyPEM []byte             // ECDH (Nebula)
-	HostEd25519PublicKeyPEM   []byte             // EdDSA (DN API)
-	HostEd25519PrivateKey     ed25519.PrivateKey // EdDSA (DN API)
+	NebulaX25519PublicKeyPEM  []byte     // ECDH (Nebula)
+	NebulaX25519PrivateKeyPEM []byte     // ECDH (Nebula)
+	HostEd25519PublicKeyPEM   []byte     // EdDSA (DN API)
+	HostEd25519PrivateKey     PrivateKey // EdDSA (DN API)
 
 	// P256 Curve
-	NebulaP256PublicKeyPEM  []byte            // ECDH (Nebula)
-	NebulaP256PrivateKeyPEM []byte            // ECDH (Nebula)
-	HostP256PublicKeyPEM    []byte            // ECDSA (DN API)
-	HostP256PrivateKey      *ecdsa.PrivateKey // ECDSA (DN API)
+	NebulaP256PublicKeyPEM  []byte     // ECDH (Nebula)
+	NebulaP256PrivateKeyPEM []byte     // ECDH (Nebula)
+	HostP256PublicKeyPEM    []byte     // ECDSA (DN API)
+	HostP256PrivateKey      PrivateKey // ECDSA (DN API)
 }
 
 func New() (*Keys, error) {
@@ -114,15 +114,25 @@ func New() (*Keys, error) {
 		return nil, err
 	}
 
+	ed25519PrivateKeyI, err := NewPrivateKey(ed25519PrivateKey)
+	if err != nil {
+		return nil, err
+	}
+
+	ecdsaP256PrivateKeyI, err := NewPrivateKey(ecdsaP256PrivateKey)
+	if err != nil {
+		return nil, err
+	}
+
 	return &Keys{
 		NebulaX25519PublicKeyPEM:  x25519PublicKeyPEM,
 		NebulaX25519PrivateKeyPEM: x25519PrivateKeyPEM,
 		HostEd25519PublicKeyPEM:   ed25519PublicKeyPEM,
-		HostEd25519PrivateKey:     ed25519PrivateKey,
+		HostEd25519PrivateKey:     ed25519PrivateKeyI,
 		NebulaP256PublicKeyPEM:    ecdhP256PublicKeyPEM,
 		NebulaP256PrivateKeyPEM:   ecdhP256PrivateKeyPEM,
 		HostP256PublicKeyPEM:      ecdsaP256PublicKeyPEM,
-		HostP256PrivateKey:        ecdsaP256PrivateKey,
+		HostP256PrivateKey:        ecdsaP256PrivateKeyI,
 	}, nil
 }