BLS12381.hs

-- SPDX-FileCopyrightText: 2021 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

-- | Support for BLS12-381 elliptic curve.
--
-- Some general hints on the implementation can be found in this python
-- re-implementation used by Tezos for testing:
-- <https://gitlab.com/metastatedev/tezos/-/commit/f10c39e0030e6b4fdd416a62de7b80b6ffdfeacf#80b4b1585c1e6fa82f2cfaf75001c490613f22c3>.
-- And it uses this library inside: <https://github.com/ethereum/py_ecc/tree/master/py_ecc/optimized_bls12_381>.
module Tezos.Crypto.BLS12381
  ( Bls12381Fr
  , Bls12381G1
  , Bls12381G2
  , CurveObject (..)
  , fromMichelsonBytesUnsafe
  , MultiplyPoint (..)
  , DeserializationError (..)
  , checkPairing

    -- * Playground
  , readFromHexUnsafe
  , generateFrom
  , g1One
  , g2One
  ) where

import Prelude hiding (negate, one)
import qualified Prelude

import Control.Exception (assert)
import Control.Lens (each, toListOf)
import Control.Monad.Random (MonadRandom, evalRand, getRandom, mkStdGen)
import Data.Bits (bit, complement, setBit, testBit, (.&.))
import qualified Data.ByteString as BS
import qualified Data.Curve as C
import qualified Data.Curve.Weierstrass as CW
import qualified Data.Curve.Weierstrass.BLS12381 as CW.BLS
import qualified Data.Field.Galois as GF
import qualified Data.Pairing.BLS12381 as BLS
import Fmt (Buildable(..), pretty, (+|), (|+))
import Named (arg, type (:!), (!))
import Text.Hex (decodeHex, encodeHex)
import Util.Instances ()
import Util.Named ()
import Util.Num

-- | Methods common for all BLS12-381 primitives.
class CurveObject a where
  -- | Representation of @0@, aka additive identity.
  zero :: a

  -- | Negate a value.
  negate :: a -> a

  -- | Add up two values.
  add :: a -> a -> a

  -- | Generate a random value.
  generate :: MonadRandom m => m a

  -- | Read a value from Michelson's bytes form.
  --
  -- Michelson tends to represent all BLS12-381 types in bytes form,
  -- some special types also allow other forms.
  fromMichelsonBytes :: ByteString -> Either DeserializationError a

  -- | Produce Michelson's bytes representation.
  toMichelsonBytes :: a -> ByteString

-- | Generate a random value from given seed.
generateFrom :: (CurveObject a) => Int -> a
generateFrom = evalRand generate . mkStdGen

-- | Read a value from Michelson's bytes form assuming that it is correct.
fromMichelsonBytesUnsafe :: (CurveObject a, HasCallStack) => ByteString -> a
fromMichelsonBytesUnsafe = either (error . pretty) id . fromMichelsonBytes

-- | Reads an object from hex string.
--
-- To be used only in playground and tests.
readFromHexUnsafe :: (CurveObject a, HasCallStack) => String -> a
readFromHexUnsafe hex =
  let bs = decodeHex (toText hex) ?: error "bad hex"
  in fromMichelsonBytesUnsafe bs

-- | Multiplication operations on BLS12-381 objects.
class MultiplyPoint scalar point where
  -- | Multiply point value by scalar value.
  multiply :: scalar -> point -> point

-- | G1 point on the curve.
newtype Bls12381G1 = Bls12381G1 { unBls12381G1 :: BLS.G1' }
  deriving stock (Show, Eq)
  deriving newtype (NFData)

instance CurveObject Bls12381G1 where
  zero = Bls12381G1 C.id
  negate (Bls12381G1 v) = Bls12381G1 (C.inv v)
  add (Bls12381G1 a) (Bls12381G1 b) = Bls12381G1 (C.add a b)
  generate = Bls12381G1 <$> C.rnd
  fromMichelsonBytes =
    let bsToCoord = toPrime . fromBigEndian
    in fmap Bls12381G1 . parseJA2WAPoint g1CoordLen bsToCoord
  toMichelsonBytes =
    let coordToBs = toBigEndian g1CoordLen . fromPrime
    in convertWA2JAPoint g1CoordLen coordToBs . unBls12381G1

instance MultiplyPoint Integer Bls12381G1 where
  multiply s (Bls12381G1 p) = Bls12381G1 (C.mul' p s)

-- | G2 point on the curve.
newtype Bls12381G2 = Bls12381G2 { unBls12381G2 :: BLS.G2' }
  deriving stock (Show, Eq)
  deriving newtype (NFData)

instance CurveObject Bls12381G2 where
  zero = Bls12381G2 C.id
  negate (Bls12381G2 v) = Bls12381G2 (C.inv v)
  add (Bls12381G2 a) (Bls12381G2 b) = Bls12381G2 (C.add a b)
  generate = Bls12381G2 <$> C.rnd
  fromMichelsonBytes =
    let fromBsPair = map fromBigEndian . toListOf each . BS.splitAt (g2CoordLen `div` 2)
        bsToCoord = GF.toE . reverse . map toPrime . fromBsPair
    in fmap Bls12381G2 . parseJA2WAPoint g2CoordLen bsToCoord
  toMichelsonBytes =
    let toBsPair = foldMap (toBigEndian $ g2CoordLen `div` 2)
        coordToBs = toBsPair . map fromPrime . reverse . GF.fromE
    in convertWA2JAPoint g1CoordLen coordToBs . unBls12381G2

instance MultiplyPoint Integer Bls12381G2 where
  multiply s (Bls12381G2 p) = Bls12381G2 (C.mul' p s)

-- | An element of an algebraic number field (scalar), used for multiplying
-- 'Bls12381G1' and 'Bls12381G2'.
newtype Bls12381Fr = Bls12381Fr { unBls12381Fr :: BLS.Fr }
  deriving stock (Show, Eq, Ord)
  deriving newtype (Num, Enum, Bounded, Real, Integral, Fractional, NFData)

instance CurveObject Bls12381Fr where
  zero = Bls12381Fr 0
  negate = Prelude.negate
  add = (+)
  generate = Bls12381Fr <$> getRandom
  fromMichelsonBytes bs =
    if length bs > frLen
    then Left $ TooLargeLength ! #limit frLen ! #given (length bs)
    else
      let num = fromLittleEndian bs
      in fromIntegralChecked num
           & first (\_ -> ValueOutsideOfField $ toInteger num)
  toMichelsonBytes = toLittleEndian frLen . fromPrime . unBls12381Fr

instance MultiplyPoint Bls12381Fr Bls12381G1 where
  multiply (Bls12381Fr s) (Bls12381G1 p) = Bls12381G1 (C.mul p s)

instance MultiplyPoint Bls12381Fr Bls12381G2 where
  multiply (Bls12381Fr s) (Bls12381G2 p) = Bls12381G2 (C.mul p s)

-- | Checks that product of pairings of points in the list is equal to 1 in
-- Fq12 field.
checkPairing :: [(Bls12381G1, Bls12381G2)] -> Bool
checkPairing pairs =
  -- Some hints for implementation of this function:
  -- https://gitlab.com/metastatedev/tezos/-/commit/bb2cda17d48a52ce854e027f0222a0463e0e66f0#af97cb649204420968454a94e7bfaa6a6e27195a_1285_1330
  -- https://gitlab.com/metastatedev/tezos/-/commit/f10c39e0030e6b4fdd416a62de7b80b6ffdfeacf#80b4b1585c1e6fa82f2cfaf75001c490613f22c3_0_172

  -- Monoid instance on GT' has the desired multiplicative semantics
  foldMap pairing pairs == mempty
  where
    pairing :: (Bls12381G1, Bls12381G2) -> BLS.GT'
    pairing (Bls12381G1 g1, Bls12381G2 g2) =
      BLS.finalExponentiationBLS12 BLS.parameterHex
        (BLS.millerAlgorithmBLS12 BLS.parameterBin g1 g2)

----------------------------------------------------------------------------
-- Serialization helpers
----------------------------------------------------------------------------

-- | All kinds of errors that can occur when reading a Michelson value.
data DeserializationError
  = CompressedFormIsNotSupported
  | UnexpectedLength ("expected" :! Int) ("given" :! Int)
  | TooLargeLength ("limit" :! Int) ("given" :! Int)
  | ValueOutsideOfField Integer
  | PointNotOnCurve ByteString
  deriving stock (Show, Eq, Generic)
  deriving anyclass (NFData)

instance Buildable DeserializationError where
  build = \case
    CompressedFormIsNotSupported ->
      "Compressed form of BLS12-381 point is not supported by Tezos"
    UnexpectedLength (arg #expected -> expected) (arg #given -> given) ->
      "Unexpected length of BLS12-381 primitive: \
      \expected " +| expected |+ ", but given " +| given |+ ""
    TooLargeLength (arg #limit -> limit) (arg #given -> given) ->
      "Too large length of BLS12-381 primitive: \
      \limit is " +| limit |+ ", but given " +| given |+ ""
    ValueOutsideOfField v ->
      "Value is too large for the given field of values: " +| v |+ ""
    PointNotOnCurve bs ->
      "Point is not on curve: 0x" +| encodeHex bs |+ ""

{- Note on serialization:

In BLS12-381, "381" stands for the number of bits necessary to represent a
coordinate of a point on a curve, i.e. we have to use 48 bytes, getting 3 extra
bits. In the BLS12-381 library used by Tezos, those bits are exploited to carry
some meta information.

More info can be found here:
<https://github.com/zkcrypto/pairing/blob/307aa1f29dccaed09abe774d2027cad57fc5d0b4/src/bls12_381/README.md#serialization>.

So Fr is just a scalar (but a pretty big one), represented in little-endian
as said in the Michelson docs.

G1 and G2 represent a point on curve and have the following form:

      X coordinate         Y coordinate
 |!___________________|____________________|
   \  (big-endian)         (big-endian)
    \
     `- bits with meta info

Generally, various coordinate systems may be used to represent a point on curve,
and the library picked by Tezos uses Jacobian coordinates, probably assuming
that the third @Z@ coordinate is always @1@.
Note that we use 'Data.Pairing.BLS12381' which by default picks Weierstrass
coordinates (it has two, not three coordinates, and an "infinity point" which
is kept as a special case), but it also provides methods for converting between
different coordinate systems.

Coordinates take a different amount of space in G1 and G2:
* In G1 both X and Y are from a so-called "Fr" field, where numbers take
  48 bytes (without the leading 3 bits).
* In G2 both coordinates are from "Fr2" field which is a two-dimensional field
  over "Fr", i.e. X and Y themselves contain two 48-byte coordinates each.

Note that this is correct for the "uncompressed" form, and there is a different
"compressed" form that, fortunately, seems to be not supported by Michelson.
They initially planned to have @COMPRESS@ and @UNCOMPRESS@ instructions, perhaps
for manual conversions, but those instructions didn't appear in Edonet
eventually.

-}

-- | A helper datatype for representing points in raw bytes form.
data RawPoint
  = Infinity
    -- ^ Point at infinity.
  | RawPoint ByteString
    -- ^ Bytes representing the payload.

-- | Given the Michelson representation of a point, interpret flags
-- and return the bare payload.
--
-- This assumes that a proper number of bytes is provided.
parsePointFlags
  :: HasCallStack
  => ByteString -> Either DeserializationError RawPoint
parsePointFlags bsFull =
  case BS.uncons bsFull of
    Nothing -> error "Empty byte sequence"
    Just (b, bs)
      | b `testBit` compressionBit ->
          Left CompressedFormIsNotSupported
      | b `testBit` infinityBit ->
          return Infinity
      | otherwise -> do
          let
            b' = b .&. complement
              (sum $ map bit [compressionBit, infinityBit, flag3Bit])
          return $ RawPoint (BS.cons b' bs)

-- | Fill a point in raw bytes form with the necessary flags.
fillPointFlags :: HasCallStack => Int -> RawPoint -> ByteString
fillPointFlags 0 = error "Coordinates are unexpectedly empty"
fillPointFlags len = \case
  Infinity -> BS.cons (0 `setBit` infinityBit) (BS.replicate (len - 1) 0)
  RawPoint bs -> bs

-- | Get a bytestring containing coordinates of a point and split it,
-- checking that each coordinate occupies the given number of bytes.
splitUncompressedPoint
  :: Int -> ByteString -> Either DeserializationError (ByteString, ByteString)
splitUncompressedPoint coordLen bs
  | length bs /= coordLen * 2 =
      Left $ UnexpectedLength ! #expected (coordLen * 2) ! #given (length bs)
  | otherwise =
      Right $ BS.splitAt coordLen bs

-- | Parse a point in Weierstrass form and Affine coordinates,
-- assuming that in the provided bytestring the point is given in Jacobian
-- coordinates (the library used by Tezos operates in Jacobian coordinates).
parseJA2WAPoint
  :: ( CW.BLS.WJCurve CW.BLS.BLS12381 fq BLS.Fr
     , CW.BLS.WACurve CW.BLS.BLS12381 fq BLS.Fr
     )
  => Int
  -> (ByteString -> fq)
  -> ByteString
  -> Either DeserializationError (CW.BLS.WAPoint CW.BLS.BLS12381 fq BLS.Fr)
parseJA2WAPoint coordLen toCoord full = do
  (xRawWithFlags, yRaw) <- splitUncompressedPoint coordLen full
  xRawPoint <- parsePointFlags xRawWithFlags
  case xRawPoint of
    Infinity -> return CW.O
    RawPoint xRaw ->
      let point = C.toA $ CW.J (toCoord xRaw) (toCoord yRaw) 1
      in if C.def point
         then return point
         else Left $ PointNotOnCurve full

-- | Turn a Weierstrass Affine point into Jacobian coordinates and represent
-- those as bytes.
convertWA2JAPoint
  :: ( CW.BLS.WJCurve CW.BLS.BLS12381 fq BLS.Fr
     , CW.BLS.WACurve CW.BLS.BLS12381 fq BLS.Fr
     )
  => Int
  -> (fq -> ByteString)
  -> CW.BLS.WAPoint CW.BLS.BLS12381 fq BLS.Fr
  -> ByteString
convertWA2JAPoint coordLen toRawCoord point =
  let
    rawPoint = case point of
      CW.O -> Infinity
      p@CW.A{} -> RawPoint $
        let CW.J x y z = C.fromA p
           -- Conversion from Affine coordinates used to produce an already
           -- normalized value.
           -- In case this turns out to be incorrect, probably @x / z@ and @y / z@
           -- are just what we want.
        in assert (z == 1) $
            toRawCoord x <> toRawCoord y
  in fillPointFlags (coordLen * 2) rawPoint

-- | Interpret a byte sequence as a number in big-endian.
fromBigEndian :: ByteString -> Natural
fromBigEndian bs =
  foldl' (\acc byte -> acc * 0x100 + fromIntegral @Word8 @Natural byte) 0 $
  BS.unpack bs

-- | Interpret a byte sequence as a number in little-endian.
fromLittleEndian :: ByteString -> Natural
fromLittleEndian bs =
  sum . zipWith (*) (iterate (* 0x100) 1) . map (fromIntegral @Word8 @Natural) $
  BS.unpack bs

-- | Represent a number in a big-endian byte sequence, padding the output
-- to the expected length.
--
-- We assert that the length is sufficient for representing the given number.
toBigEndian :: Int -> Natural -> ByteString
toBigEndian len num =
  BS.pack $
  let
    (remainder, bytes) = mapAccumR
      (\x _ -> second (fromIntegral @Natural @Word8) $ x `divMod` 0x100)
      num [1 .. len]
  in assert (remainder == 0) bytes

-- | Represent a number in a little-endian byte sequence, padding the output
-- to the expected length.
--
-- We assert that the length is sufficient for representing the given number.
toLittleEndian :: Int -> Natural -> ByteString
toLittleEndian len num =
  BS.pack $
  let
    (remainder, bytes) = mapAccumL
      (\x _ -> second (fromIntegral @Natural @Word8) $ x `divMod` 0x100)
      num [1 .. len]
  in assert (remainder == 0) bytes

-- | Turn a prime field element into a natural.
fromPrime :: KnownNat p => GF.Prime p -> Natural
fromPrime p =
  -- This should be safe, since 'Prime's exist in modular arithmetics,
  -- so its conversion to an integer should produce non-negative elements.
  -- In fact, 'GF.Prime' is a newtype wrapper over 'Natural', but
  -- its constructor is not exported :/
  assert (p >= 0) $
    fromIntegral p

-- | The inverse to 'fromPrime'.
toPrime :: KnownNat p => Natural -> GF.Prime p
toPrime = fromIntegral

-- Primitives' lengths
----------------------------------------------------------------------------

-- | Length of a single coordinate of a point in raw bytes form.
g1CoordLen, g2CoordLen :: Int
g1CoordLen = 48
g2CoordLen = 96  -- each coordinate is an element of two-dimensional field Fr2

-- | Length of 'Fr' in raw bytes form.
frLen :: Int
frLen = 32

-- Meta bits
----------------------------------------------------------------------------

-- | This bit designates whether the point is represented in compressed form
-- (only X coordinate), or uncompressed form (both X and Y coordinates).
compressionBit :: Int
compressionBit = 7

-- | This bit designates whether the given point is at infinity.
--
-- If so, all other bytes should be zeros.
infinityBit :: Int
infinityBit = 6

-- | This bit is set iff "this point is in compressed form /and/ it is not the
-- point at infinity /and/ its y-coordinate is the lexicographically largest of
-- the two associated with the encoded x-coordinate".
--
-- Fortunatelly, this flag seems to be not relevant for us at the moment.
flag3Bit :: Int
flag3Bit = 5

----------------------------------------------------------------------------
-- Other constants
----------------------------------------------------------------------------

-- | @1@ represented in G1 - as the libraries used by Tezos see it.
--
-- Taken from here: <https://github.com/ethereum/py_ecc/blob/3f644b4c07c8270b8fbe989eb799766aca66face/py_ecc/optimized_bls12_381/optimized_curve.py#L34>.
g1One :: Bls12381G1
g1One = Bls12381G1 $ CW.toA $ CW.J
  3685416753713387016781088315183077757961620795782546409894578378688607592378376318836054947676345821548104185464507
  1339506544944476473020471379941921221584933875938349620426543736416511423956333506472724655353366534992391756441569
  1

-- | @1@ represented in G2.
g2One :: Bls12381G2
g2One = Bls12381G2 $ CW.toA $ CW.J
  (GF.toE
    [ 352701069587466618187139116011060144890029952792775240219908644239793785735715026873347600343865175952761926303160
    , 3059144344244213709971259814753781636986470325476647558659373206291635324768958432433509563104347017837885763365758
    ])
  (GF.toE
    [ 1985150602287291935568054521177171638300868978215655730859378665066344726373823718423869104263333984641494340347905
    , 927553665492332455747201965776037880757740193453592970025027978793976877002675564980949289727957565575433344219582
    ])
  (GF.toE [1, 0])