From c6e775163866d6ea5233eb8ec8530a9122101ebd Mon Sep 17 00:00:00 2001
From: Savely Krasovsky <savely@krasovs.ky>
Date: Mon, 4 Dec 2023 21:29:05 +0100
Subject: [PATCH] fix: additional security checks added

---
 privatekey.go | 8 ++++++++
 publickey.go  | 6 +++++-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/privatekey.go b/privatekey.go
index 38e8b33..a2db9a2 100644
--- a/privatekey.go
+++ b/privatekey.go
@@ -77,6 +77,10 @@ func (k *PrivateKey) Encapsulate(pub *PublicKey) ([]byte, error) {
 		return nil, fmt.Errorf("public key is empty")
 	}
 
+	if !k.Curve.IsOnCurve(pub.X, pub.Y) {
+		return nil, fmt.Errorf("invalid public key")
+	}
+
 	var secret bytes.Buffer
 	secret.Write(k.PublicKey.Bytes(false))
 
@@ -98,6 +102,10 @@ func (k *PrivateKey) ECDH(pub *PublicKey) ([]byte, error) {
 		return nil, fmt.Errorf("public key is empty")
 	}
 
+	if !k.Curve.IsOnCurve(pub.X, pub.Y) {
+		return nil, fmt.Errorf("invalid public key")
+	}
+
 	// Shared secret generation
 	sx, sy := pub.Curve.ScalarMult(pub.X, pub.Y, k.D.Bytes())
 
diff --git a/publickey.go b/publickey.go
index dbd1adb..da1ffaf 100644
--- a/publickey.go
+++ b/publickey.go
@@ -137,8 +137,12 @@ func (k *PublicKey) Hex(compressed bool) string {
 // Decapsulate decapsulates key by using Key Encapsulation Mechanism and returns symmetric key;
 // can be safely used as encryption key
 func (k *PublicKey) Decapsulate(priv *PrivateKey) ([]byte, error) {
+	if !k.Curve.IsOnCurve(k.X, k.Y) {
+		return nil, fmt.Errorf("invalid public key")
+	}
+
 	if priv == nil {
-		return nil, fmt.Errorf("public key is empty")
+		return nil, fmt.Errorf("private key is empty")
 	}
 
 	var secret bytes.Buffer