Use ecdh package for more (un)marshal operations

smallstep · Jul 9, 2024 · 3e863b2 · 3e863b2
1 parent f9dd5e7
commit 3e863b2
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 12 deletions.
diff --git a/pemutil/ssh.go b/pemutil/ssh.go
@@ -10,7 +10,6 @@ import (
 	"crypto/cipher"
 	"crypto/ecdsa"
 	"crypto/ed25519"
-	"crypto/elliptic"
 	"crypto/rand"
 	"crypto/rsa"
 	"encoding/binary"
@@ -188,15 +187,18 @@ func SerializeOpenSSHPrivateKey(key crypto.PrivateKey, opts ...Options) (*pem.Bl
 			return nil, errors.Errorf("error serializing key: unsupported curve %s", k.Curve.Params().Name)
 		}
 
-		pub := elliptic.Marshal(k.Curve, k.PublicKey.X, k.PublicKey.Y)
+		p, err := k.PublicKey.ECDH()
+		if err != nil {
+			return nil, errors.Wrapf(err, "failed converting *ecdsa.PublicKey to *ecdh.PublicKey")
+		}
 
 		// Marshal public key.
 		pubKey := struct {
 			KeyType string
 			Curve   string
 			Pub     []byte
 		}{
-			keyType, curve, pub,
+			keyType, curve, p.Bytes(),
 		}
 		w.PubKey = ssh.Marshal(pubKey)
 
@@ -207,7 +209,7 @@ func SerializeOpenSSHPrivateKey(key crypto.PrivateKey, opts ...Options) (*pem.Bl
 			D       *big.Int
 			Comment string
 		}{
-			curve, pub, k.D,
+			curve, p.Bytes(), k.D,
 			ctx.comment,
 		}
 		pk1.Keytype = keyType

diff --git a/sshutil/sshutil.go b/sshutil/sshutil.go
@@ -3,11 +3,13 @@ package sshutil
 import (
 	"crypto"
 	"crypto/dsa" //nolint:staticcheck // support for DSA fingerprints
+	"crypto/ecdh"
 	"crypto/ecdsa"
 	"crypto/ed25519"
 	"crypto/elliptic"
 	"crypto/rsa"
 	"fmt"
+	"math/big"
 
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/crypto/ssh/agent"
@@ -48,13 +50,17 @@ func cryptoSKPublicKey(pub ssh.PublicKey) (crypto.PublicKey, error) {
 		if err := ssh.Unmarshal(pub.Marshal(), &w); err != nil {
 			return nil, err
 		}
-		key := new(ecdsa.PublicKey)
-		key.Curve = elliptic.P256()
-		key.X, key.Y = elliptic.Unmarshal(key.Curve, w.Key)
-		if key.X == nil || key.Y == nil {
-			return nil, fmt.Errorf("invalid curve point")
+
+		p, err := ecdh.P256().NewPublicKey(w.Key)
+		if err != nil {
+			return nil, fmt.Errorf("failed decoding ECDSA key: %w", err)
 		}
-		return key, nil
+
+		return &ecdsa.PublicKey{
+			Curve: elliptic.P256(),
+			X:     big.NewInt(0).SetBytes(p.Bytes()[1:33]),
+			Y:     big.NewInt(0).SetBytes(p.Bytes()[33:]),
+		}, nil
 	case "sk-ssh-ed25519@openssh.com":
 		var w struct {
 			Name        string

diff --git a/sshutil/sshutil_test.go b/sshutil/sshutil_test.go
@@ -4,10 +4,10 @@ import (
 	"crypto"
 	"crypto/ecdsa"
 	"crypto/ed25519"
-	"crypto/elliptic"
 	"reflect"
 	"testing"
 
+	"github.com/stretchr/testify/require"
 	"go.step.sm/crypto/keyutil"
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/crypto/ssh/agent"
@@ -36,14 +36,17 @@ func generateKey(t *testing.T, kty, crv string, size int) (crypto.PublicKey, ssh
 }
 
 func generateFakeSKKey(t *testing.T, pub crypto.PublicKey) ssh.PublicKey {
+	t.Helper()
 	switch k := pub.(type) {
 	case *ecdsa.PublicKey:
+		p, err := k.ECDH()
+		require.NoError(t, err)
 		w := struct {
 			Name        string
 			ID          string
 			Key         []byte
 			Application string
-		}{"sk-ecdsa-sha2-nistp256@openssh.com", "nistp256", elliptic.Marshal(k.Curve, k.X, k.Y), "ssh"}
+		}{"sk-ecdsa-sha2-nistp256@openssh.com", "nistp256", p.Bytes(), "ssh"}
 		return &skKey{
 			typ:   "sk-ecdsa-sha2-nistp256@openssh.com",
 			bytes: ssh.Marshal(w),