add extra error check

hppk.go

@@ -45,8 +45,8 @@ type PrivateKey struct {
 
 // PublicKey represents a public key in the HPPK protocol.
 type PublicKey struct {
-	P [][]*big.Int // Coefficient matrix of the polynomial P(x)
-	Q [][]*big.Int
+	Pm [][]*big.Int // Coefficient matrix of the polynomial P(x)
+	Qm [][]*big.Int
 }
 
 // Signature represents a digital signature in the HPPK protocol.
@@ -187,8 +187,8 @@ RETRY:
 		H0:    h0,
 		H1:    h1,
 		PublicKey: PublicKey{
-			P: Pm,
-			Q: Qm,
+			Pm: Pm,
+			Qm: Qm,
 		},
 	}, nil
 }
@@ -221,16 +221,20 @@ func encrypt(pub *PublicKey, msg []byte, prime *big.Int) (kem *KEM, err error) {
 	}
 
 	// Ensure fields in the public key are valid
-	if pub.P == nil || pub.Q == nil {
+	if len(pub.Pm) == 0 || len(pub.Qm) == 0 {
 		return nil, errors.New(ERR_MSG_INVALID_PUBKEY)
 	}
 
-	if len(pub.P) != len(pub.Q) {
+	if len(pub.Pm) != len(pub.Qm) {
 		return nil, errors.New(ERR_MSG_INVALID_PUBKEY)
 	}
 
-	for i := 0; i < len(pub.P); i++ {
-		if pub.P[i] == nil || pub.Q[i] == nil {
+	for i := 0; i < len(pub.Pm); i++ {
+		if pub.Pm[i] == nil || pub.Qm[i] == nil {
+			return nil, errors.New(ERR_MSG_INVALID_PUBKEY)
+		}
+
+		if len(pub.Pm[i]) != len(pub.Qm[i]) {
 			return nil, errors.New(ERR_MSG_INVALID_PUBKEY)
 		}
 	}
@@ -239,7 +243,7 @@ func encrypt(pub *PublicKey, msg []byte, prime *big.Int) (kem *KEM, err error) {
 	P := new(big.Int)
 	Q := new(big.Int)
 
-	for c := 0; c < len(pub.P); c++ {
+	for c := 0; c < len(pub.Pm); c++ {
 		// Generate a random noise
 		noise, err := rand.Int(rand.Reader, prime)
 		if err != nil {
@@ -249,12 +253,12 @@ func encrypt(pub *PublicKey, msg []byte, prime *big.Int) (kem *KEM, err error) {
 		// Initialize Si with the secret message
 		Si := big.NewInt(1)
 		t := new(big.Int)
-		for i := 0; i < len(pub.P[0]); i++ {
+		for i := 0; i < len(pub.Pm[0]); i++ {
 			noised := new(big.Int).Mul(noise, Si)
 			noised.Mod(noised, prime)
 
-			P.Add(P, t.Mul(Si, pub.P[c][i]))
-			Q.Add(Q, t.Mul(Si, pub.Q[c][i]))
+			P.Add(P, t.Mul(Si, pub.Pm[c][i]))
+			Q.Add(Q, t.Mul(Si, pub.Qm[c][i]))
 
 			// Si = secret^i
 			Si.Mul(Si, secret)
@@ -393,13 +397,13 @@ func (priv *PrivateKey) Sign(digest []byte) (sign *Signature, err error) {
 	R := new(big.Int).Exp(big.NewInt(2), big.NewInt(int64(K)), nil)
 
 	// Initiate V, U
-	V := make([]*big.Int, len(priv.P[0]))
-	U := make([]*big.Int, len(priv.Q[0]))
+	V := make([]*big.Int, len(priv.Pm[0]))
+	U := make([]*big.Int, len(priv.Qm[0]))
 	for i := 0; i < len(V); i++ {
-		V[i] = new(big.Int).Mul(priv.Q[0][i], R)
+		V[i] = new(big.Int).Mul(priv.Qm[0][i], R)
 		V[i].Quo(V[i], priv.S2)
 
-		U[i] = new(big.Int).Mul(priv.P[0][i], R)
+		U[i] = new(big.Int).Mul(priv.Pm[0][i], R)
 		U[i].Quo(U[i], priv.S1)
 	}
 
@@ -423,7 +427,7 @@ func (priv *PrivateKey) Public() *PublicKey {
 
 // Order returns the polynomial order of the private key.
 func (priv *PrivateKey) Order() int {
-	return len(priv.PublicKey.P[0]) - 2
+	return len(priv.PublicKey.Pm[0]) - 2
 }
 
 // VerifySignature verifies the signature of the message digest using the public key and given prime
@@ -451,16 +455,20 @@ func verifySignature(sig *Signature, digest []byte, pub *PublicKey, prime *big.I
 		return false
 	}
 
-	if len(pub.P) == 0 || len(pub.Q) == 0 {
+	if len(pub.Pm) == 0 || len(pub.Qm) == 0 {
 		return false
 	}
 
-	if len(pub.P) != len(pub.Q) {
+	if len(pub.Pm) != len(pub.Qm) {
 		return false
 	}
 
-	for i := 0; i < len(pub.P); i++ {
-		if pub.P[i] == nil || pub.Q[i] == nil {
+	for i := 0; i < len(pub.Pm); i++ {
+		if pub.Pm[i] == nil || pub.Qm[i] == nil {
+			return false
+		}
+
+		if len(pub.Pm[i]) != len(pub.Qm[i]) {
 			return false
 		}
 	}
@@ -469,10 +477,10 @@ func verifySignature(sig *Signature, digest []byte, pub *PublicKey, prime *big.I
 	Q := make([]*big.Int, len(sig.U))
 	P := make([]*big.Int, len(sig.V))
 	for i := 0; i < len(Q); i++ {
-		Q[i] = new(big.Int).Mul(pub.Q[0][i], sig.Beta)
+		Q[i] = new(big.Int).Mul(pub.Qm[0][i], sig.Beta)
 		Q[i].Mod(Q[i], prime)
 
-		P[i] = new(big.Int).Mul(pub.P[0][i], sig.Beta)
+		P[i] = new(big.Int).Mul(pub.Pm[0][i], sig.Beta)
 		P[i].Mod(P[i], prime)
 	}