fix: additional security checks added

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())
 

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