Plaintext space check

According to [DJN 10], the plaintext space is `Z_n`. We had this
information in the `Encrypt` and `EncryptWithR` documentation but
sometimes, when encrypting result of more complex equations, it's easy
to miss the fact the result is out of the allowed range.

Here we have just a safety check making sure that if plaintext is out of
allowed space, we report an error.

paillier.go

@@ -15,8 +15,8 @@ func (pk *PublicKey) GetNSquare() *big.Int {
 }
 
 // EncryptWithR encrypts a plaintext into a cypher one with random `r` specified
-// in the argument. The plain text must be smaller that N and bigger or equal
-// than zero. `r` is the randomness used to encrypt the plaintext. `r` must be
+// in the argument. The plain text must be smaller that N and bigger than or
+// equal zero. `r` is the randomness used to encrypt the plaintext. `r` must be
 // a random element from a multiplicative group of integers modulo N.
 //
 // If you don't need to use the specific `r`, you should use the `Encrypt`
@@ -28,6 +28,14 @@ func (pk *PublicKey) GetNSquare() *big.Int {
 //
 // See [KL 08] construction 11.32, page 414.
 func (pk *PublicKey) EncryptWithR(m *big.Int, r *big.Int) (*Cypher, error) {
+	if m.Cmp(ZERO) == -1 || m.Cmp(pk.N) != -1 { // m < 0 || m >= N  ?
+		return nil, fmt.Errorf(
+			"%v is out of allowed plaintext space [0, %v)",
+			m,
+			pk.N,
+		)
+	}
+
 	nSquare := pk.GetNSquare()
 
 	// g is _always_ equal n+1
@@ -40,7 +48,7 @@ func (pk *PublicKey) EncryptWithR(m *big.Int, r *big.Int) (*Cypher, error) {
 }
 
 // Encrypt a plaintext into a cypher one. The plain text must be smaller that
-// N and bigger or equal than zero. random is usually rand.Reader from the
+// N and bigger than or equal zero. random is usually rand.Reader from the
 // package crypto/rand.
 //
 // m - plaintext to encrypt

paillier_test.go

@@ -2,7 +2,9 @@ package paillier
 
 import (
 	"crypto/rand"
+	"errors"
 	"math/big"
+	"reflect"
 	"testing"
 )
 
@@ -63,6 +65,65 @@ func TestEncryptDecryptSmall(t *testing.T) {
 	}
 }
 
+func TestCheckPlaintextSpace(t *testing.T) {
+	p := big.NewInt(13)
+	q := big.NewInt(11)
+
+	// N = pq = 143 so the plaintext space is [0, 143)
+	privateKey := CreatePrivateKey(p, q)
+
+	var tests = map[string]struct {
+		plaintext     *big.Int
+		expectedError error
+	}{
+		"plaintext less than 0": {
+			plaintext:     big.NewInt(-1),
+			expectedError: errors.New("-1 is out of allowed plaintext space [0, 143)"),
+		},
+		"plaintext equal 0": {
+			plaintext: big.NewInt(0),
+		},
+		"plaintext equal 1": {
+			plaintext: big.NewInt(1),
+		},
+		"plaintext equal 142": {
+			plaintext: big.NewInt(142),
+		},
+		"plaintext equal 143": {
+			plaintext:     big.NewInt(143),
+			expectedError: errors.New("143 is out of allowed plaintext space [0, 143)"),
+		},
+		"plaintext equal 144": {
+			plaintext:     big.NewInt(144),
+			expectedError: errors.New("144 is out of allowed plaintext space [0, 143)"),
+		},
+	}
+
+	for testName, test := range tests {
+		t.Run(testName, func(t *testing.T) {
+			cypher, err := privateKey.Encrypt(test.plaintext, rand.Reader)
+			if !reflect.DeepEqual(err, test.expectedError) {
+				t.Errorf(
+					"Unexpected error\nExpected: %v\nActual: %v",
+					test.expectedError,
+					err,
+				)
+			}
+
+			if test.expectedError == nil {
+				decrypted := privateKey.Decrypt(cypher)
+				if test.plaintext.Cmp(decrypted) != 0 {
+					t.Errorf(
+						"Unexpected decryption\nExpected: %v\nActual: %v",
+						test.plaintext,
+						decrypted,
+					)
+				}
+			}
+		})
+	}
+}
+
 func TestAddCyphers(t *testing.T) {
 	privateKey := CreatePrivateKey(big.NewInt(17), big.NewInt(13))