From 47d79f75766c88ddab688f41b2c8b014af08f177 Mon Sep 17 00:00:00 2001
From: Alejandro Mery <amery@jpi.io>
Date: Wed, 23 Aug 2023 00:29:15 +0000
Subject: [PATCH] wireguard: introduce KeyPair.Validate()

it will also set the PublicKey field is empty

Signed-off-by: Alejandro Mery <amery@jpi.io>
---
 pkg/wireguard/keys.go | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/pkg/wireguard/keys.go b/pkg/wireguard/keys.go
index 100eb02..c18cf3c 100644
--- a/pkg/wireguard/keys.go
+++ b/pkg/wireguard/keys.go
@@ -19,6 +19,10 @@ const (
 var (
 	// ErrInvalidKeySize indicates the key size is wrong
 	ErrInvalidKeySize = errors.New("invalid key size")
+	// ErrInvalidPrivateKey indicates the private key is invalid
+	ErrInvalidPrivateKey = errors.New("invalid private key")
+	// ErrInvalidPublicKey indicates the public key is invalid
+	ErrInvalidPublicKey = errors.New("invalid public key")
 )
 
 type (
@@ -127,6 +131,32 @@ type KeyPair struct {
 	PublicKey  PublicKey
 }
 
+// Validate checks the PublicKey matches the PrivateKey,
+// and sets the PublicKey if missing
+func (kp *KeyPair) Validate() error {
+	keyLen := len(kp.PrivateKey)
+	pubLen := len(kp.PublicKey)
+
+	switch {
+	case keyLen != PrivateKeySize:
+		// bad private key
+		return ErrInvalidPrivateKey
+	case pubLen == 0:
+		// no public key, set it
+		kp.PublicKey = kp.PrivateKey.Public()
+		return nil
+	case pubLen != PublicKeySize:
+		// bad public key
+		return ErrInvalidPublicKey
+	case !kp.PrivateKey.Public().Equal(kp.PublicKey):
+		// wrong public key
+		return ErrInvalidPublicKey
+	default:
+		// correct public key
+		return nil
+	}
+}
+
 // NewKeyPair creates a new KeyPair for Wireguard
 func NewKeyPair() (*KeyPair, error) {
 	key, err := NewPrivateKey()