Implement unmarshal for compressed format for crypto/elliptic

First byte can be 0x2 or 0x3 (compressed form).
To unmarshal used formula y² = x³ - 3x + b.
Reuse code from `IsOnCurve`.

Closes golang#34105 issue

Author: im-kulikov

src/crypto/elliptic/elliptic.go

@@ -52,11 +52,8 @@ func (curve *CurveParams) Params() *CurveParams {
 	return curve
 }
 
-func (curve *CurveParams) IsOnCurve(x, y *big.Int) bool {
-	// y² = x³ - 3x + b
-	y2 := new(big.Int).Mul(y, y)
-	y2.Mod(y2, curve.P)
-
+// polynomial returns x³ - 3x + b
+func (curve *CurveParams) polynomial(x *big.Int) *big.Int {
 	x3 := new(big.Int).Mul(x, x)
 	x3.Mul(x3, x)
 
@@ -67,7 +64,15 @@ func (curve *CurveParams) IsOnCurve(x, y *big.Int) bool {
 	x3.Add(x3, curve.B)
 	x3.Mod(x3, curve.P)
 
-	return x3.Cmp(y2) == 0
+	return x3
+}
+
+func (curve *CurveParams) IsOnCurve(x, y *big.Int) bool {
+	// y² = x³ - 3x + b
+	y2 := new(big.Int).Mul(y, y)
+	y2.Mod(y2, curve.P)
+
+	return curve.polynomial(x).Cmp(y2) == 0
 }
 
 // zForAffine returns a Jacobian Z value for the affine point (x, y). If x and
@@ -321,21 +326,39 @@ func Marshal(curve Curve, x, y *big.Int) []byte {
 // On error, x = nil.
 func Unmarshal(curve Curve, data []byte) (x, y *big.Int) {
 	byteLen := (curve.Params().BitSize + 7) >> 3
-	if len(data) != 1+2*byteLen {
-		return
-	}
-	if data[0] != 4 { // uncompressed form
-		return
-	}
-	p := curve.Params().P
-	x = new(big.Int).SetBytes(data[1 : 1+byteLen])
-	y = new(big.Int).SetBytes(data[1+byteLen:])
-	if x.Cmp(p) >= 0 || y.Cmp(p) >= 0 {
-		return nil, nil
-	}
-	if !curve.IsOnCurve(x, y) {
-		return nil, nil
+
+	switch data[0] {
+	case 2, 3: // compressed form
+		if len(data) != 1+byteLen {
+			return
+		}
+
+		x = new(big.Int).SetBytes(data[1:])
+		x3 := curve.Params().polynomial(x)
+		y = new(big.Int).ModSqrt(x3, curve.Params().P)
+
+		// Jacobi(x, p) == -1
+		if y == nil {
+			return nil, nil
+		} else if y.Bit(0) != uint(data[0]&0x1) {
+			y.Neg(y)
+			y.Mod(y, curve.Params().P)
+		}
+	case 4: // uncompressed form
+		if len(data) != 1+2*byteLen {
+			return
+		}
+		p := curve.Params().P
+		x = new(big.Int).SetBytes(data[1 : 1+byteLen])
+		y = new(big.Int).SetBytes(data[1+byteLen:])
+		if x.Cmp(p) >= 0 || y.Cmp(p) >= 0 {
+			return nil, nil
+		}
+		if !curve.IsOnCurve(x, y) {
+			return nil, nil
+		}
 	}
+
 	return
 }
 

src/crypto/elliptic/elliptic_test.go

@@ -628,3 +628,26 @@ func TestUnmarshalToLargeCoordinates(t *testing.T) {
 		t.Errorf("Unmarshal accepts invalid Y coordinate")
 	}
 }
+
+// See https://golang.org/issues/34105
+func TestUnmarshalCompressed(t *testing.T) {
+	p256 := P256()
+	_, x, y, err := GenerateKey(p256, rand.Reader)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	byteLen := (p256.Params().BitSize + 7) >> 3
+	compressed := make([]byte, 1+byteLen)
+	compressed[0] = 3
+	if y.Bit(0) == 0 {
+		compressed[0] = 2
+	}
+	copy(compressed[1:], x.Bytes())
+	newX, newY := Unmarshal(p256, compressed)
+	if !p256.IsOnCurve(newX, newY) {
+		t.Error("P256 failed to validate a correct point")
+	} else if x.Cmp(newX) != 0 || y.Cmp(newY) != 0 {
+		t.Error("P256 failed to correctly unmarshal compressed point")
+	}
+}