verificationhelper: fix JSON serialization of keys

Signed-off-by: Sumner Evans <[email protected]>

crypto/verificationhelper/ecdhkeys.go

@@ -0,0 +1,48 @@
+package verificationhelper
+
+import (
+	"crypto/ecdh"
+	"encoding/json"
+)
+
+type ECDHPrivateKey struct {
+	*ecdh.PrivateKey
+}
+
+func (e *ECDHPrivateKey) UnmarshalJSON(data []byte) (err error) {
+	if len(data) == 0 {
+		return nil
+	}
+	var raw []byte
+	err = json.Unmarshal(data, &raw)
+	if err != nil {
+		return
+	}
+	e.PrivateKey, err = ecdh.P256().NewPrivateKey(raw)
+	return err
+}
+
+func (e ECDHPrivateKey) MarshalJSON() ([]byte, error) {
+	if e.PrivateKey == nil {
+		return json.Marshal(nil)
+	}
+	return json.Marshal(e.Bytes())
+}
+
+type ECDHPublicKey struct {
+	*ecdh.PublicKey
+}
+
+func (e *ECDHPublicKey) UnmarshalJSON(data []byte) (err error) {
+	var raw []byte
+	err = json.Unmarshal(data, &raw)
+	if err != nil {
+		return
+	}
+	e.PublicKey, err = ecdh.P256().NewPublicKey(raw)
+	return
+}
+
+func (e ECDHPublicKey) MarshalJSON() ([]byte, error) {
+	return json.Marshal(e.Bytes())
+}

crypto/verificationhelper/ecdhkeys_test.go

@@ -0,0 +1,48 @@
+package verificationhelper_test
+
+import (
+	"crypto/ecdh"
+	"crypto/rand"
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"maunium.net/go/mautrix/crypto/verificationhelper"
+)
+
+func TestECDHPrivateKey(t *testing.T) {
+	pk, err := ecdh.P256().GenerateKey(rand.Reader)
+	require.NoError(t, err)
+	private := verificationhelper.ECDHPrivateKey{pk}
+	marshalled, err := json.Marshal(private)
+	require.NoError(t, err)
+
+	assert.Len(t, marshalled, 46)
+
+	var unmarshalled verificationhelper.ECDHPrivateKey
+	err = json.Unmarshal(marshalled, &unmarshalled)
+	require.NoError(t, err)
+
+	assert.True(t, private.Equal(unmarshalled.PrivateKey))
+}
+
+func TestECDHPublicKey(t *testing.T) {
+	private, err := ecdh.P256().GenerateKey(rand.Reader)
+	require.NoError(t, err)
+
+	public := private.PublicKey()
+
+	pub := verificationhelper.ECDHPublicKey{public}
+	marshalled, err := json.Marshal(pub)
+	require.NoError(t, err)
+
+	assert.Len(t, marshalled, 90)
+
+	var unmarshalled verificationhelper.ECDHPublicKey
+	err = json.Unmarshal(marshalled, &unmarshalled)
+	require.NoError(t, err)
+
+	assert.True(t, public.Equal(unmarshalled.PublicKey))
+}

crypto/verificationhelper/verificationstore.go

@@ -2,8 +2,6 @@ package verificationhelper
 
 import (
 	"context"
-	"crypto/ecdh"
-	"encoding/json"
 	"errors"
 	"fmt"
 
@@ -53,32 +51,6 @@ func (step VerificationState) String() string {
 	}
 }
 
-type ECDHPrivateKey struct {
-	*ecdh.PrivateKey
-}
-
-func (e *ECDHPrivateKey) UnmarshalJSON(data []byte) (err error) {
-	e.PrivateKey, err = ecdh.P256().NewPrivateKey(data)
-	return
-}
-
-func (e *ECDHPrivateKey) MarshalJSON() ([]byte, error) {
-	return json.Marshal(e.Bytes())
-}
-
-type ECDHPublicKey struct {
-	*ecdh.PublicKey
-}
-
-func (e *ECDHPublicKey) UnmarshalJSON(data []byte) (err error) {
-	e.PublicKey, err = ecdh.P256().NewPublicKey(data)
-	return
-}
-
-func (e *ECDHPublicKey) MarshalJSON() ([]byte, error) {
-	return json.Marshal(e.Bytes())
-}
-
 type VerificationTransaction struct {
 	ExpirationTime jsontime.UnixMilli `json:"expiration_time"`