-- | Deterministic protobuf encoding for libp2p PublicKey/PrivateKey messages.
--
-- Hand-rolled encoding to guarantee deterministic output:
-- - Minimal varint encoding
-- - Fields in field number order (Type=1, Data=2)
-- - All fields present, no extras
--
-- Both messages share the same wire layout (per the peer-ids spec):
--
-- > message PublicKey  { required KeyType Type = 1; required bytes Data = 2; }
-- > message PrivateKey { required KeyType Type = 1; required bytes Data = 2; }
module LibP2P.Crypto.Protobuf
  ( encodePublicKey
  , decodePublicKey
  , encodePrivateKey
  , decodePrivateKey
  ) where

import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.Word (Word64, Word8)
import Numeric (showHex)
import LibP2P.Core.Varint (decodeUvarint, encodeUvarint)
import LibP2P.Crypto.Key (KeyType (..), PrivateKey (..), PublicKey (..))

-- | Protobuf KeyType enum values (per libp2p peer-ids spec: RSA=0, Ed25519=1,
-- Secp256k1=2, ECDSA=3).
keyTypeToProto :: KeyType -> Word64
keyTypeToProto :: KeyType -> Word64
keyTypeToProto KeyType
RSA = Word64
0
keyTypeToProto KeyType
Ed25519 = Word64
1
keyTypeToProto KeyType
Secp256k1 = Word64
2
keyTypeToProto KeyType
ECDSA = Word64
3

keyTypeFromProto :: Word64 -> Either String KeyType
keyTypeFromProto :: Word64 -> Either String KeyType
keyTypeFromProto Word64
0 = KeyType -> Either String KeyType
forall a b. b -> Either a b
Right KeyType
RSA
keyTypeFromProto Word64
1 = KeyType -> Either String KeyType
forall a b. b -> Either a b
Right KeyType
Ed25519
keyTypeFromProto Word64
2 = KeyType -> Either String KeyType
forall a b. b -> Either a b
Right KeyType
Secp256k1
keyTypeFromProto Word64
3 = KeyType -> Either String KeyType
forall a b. b -> Either a b
Right KeyType
ECDSA
keyTypeFromProto Word64
n = String -> Either String KeyType
forall a b. a -> Either a b
Left (String -> Either String KeyType)
-> String -> Either String KeyType
forall a b. (a -> b) -> a -> b
$ String
"unknown KeyType: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Word64 -> String
forall a. Show a => a -> String
show Word64
n

-- | Deterministic protobuf encoding of a PublicKey message.
encodePublicKey :: PublicKey -> ByteString
encodePublicKey :: PublicKey -> ByteString
encodePublicKey (PublicKey KeyType
kt ByteString
rawKey) = KeyType -> ByteString -> ByteString
encodeKeyMessage KeyType
kt ByteString
rawKey

-- | Decode a protobuf-encoded PublicKey message.
decodePublicKey :: ByteString -> Either String PublicKey
decodePublicKey :: ByteString -> Either String PublicKey
decodePublicKey ByteString
bs = (KeyType -> ByteString -> PublicKey)
-> (KeyType, ByteString) -> PublicKey
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry KeyType -> ByteString -> PublicKey
PublicKey ((KeyType, ByteString) -> PublicKey)
-> Either String (KeyType, ByteString) -> Either String PublicKey
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ByteString -> Either String (KeyType, ByteString)
decodeKeyMessage String
"decodePublicKey" ByteString
bs

-- | Deterministic protobuf encoding of a PrivateKey message. Note that
-- PrivateKey messages are never sent over the wire; this is the on-disk
-- format shared by libp2p implementations.
encodePrivateKey :: PrivateKey -> ByteString
encodePrivateKey :: PrivateKey -> ByteString
encodePrivateKey (PrivateKey KeyType
kt ByteString
rawKey) = KeyType -> ByteString -> ByteString
encodeKeyMessage KeyType
kt ByteString
rawKey

-- | Decode a protobuf-encoded PrivateKey message, e.g. a key file written
-- by another implementation. The Data field is returned as-is; use
-- 'LibP2P.Crypto.Key.keyPairFromPrivateKey' to validate it and derive the
-- public key.
decodePrivateKey :: ByteString -> Either String PrivateKey
decodePrivateKey :: ByteString -> Either String PrivateKey
decodePrivateKey ByteString
bs = (KeyType -> ByteString -> PrivateKey)
-> (KeyType, ByteString) -> PrivateKey
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry KeyType -> ByteString -> PrivateKey
PrivateKey ((KeyType, ByteString) -> PrivateKey)
-> Either String (KeyType, ByteString) -> Either String PrivateKey
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ByteString -> Either String (KeyType, ByteString)
decodeKeyMessage String
"decodePrivateKey" ByteString
bs

-- | Shared deterministic encoding of the key messages.
--
-- Layout:
--   Field 1 (Type): tag=0x08 (field 1, wire type 0=varint), value=keytype
--   Field 2 (Data): tag=0x12 (field 2, wire type 2=length-delimited), length, bytes
encodeKeyMessage :: KeyType -> ByteString -> ByteString
encodeKeyMessage :: KeyType -> ByteString -> ByteString
encodeKeyMessage KeyType
kt ByteString
rawKey =
  -- Field 1: tag 0x08 + varint value
  Word8 -> ByteString
BS.singleton Word8
0x08 ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Word64 -> ByteString
encodeUvarint (KeyType -> Word64
keyTypeToProto KeyType
kt)
    -- Field 2: tag 0x12 + varint length + raw bytes
    ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Word8 -> ByteString
BS.singleton Word8
0x12
    ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Word64 -> ByteString
encodeUvarint (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
BS.length ByteString
rawKey))
    ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
rawKey

-- | Shared decoding of the key messages; @ctx@ names the caller in errors.
decodeKeyMessage :: String -> ByteString -> Either String (KeyType, ByteString)
decodeKeyMessage :: String -> ByteString -> Either String (KeyType, ByteString)
decodeKeyMessage String
ctx ByteString
bs = do
  -- Field 1: expect tag 0x08
  (tag1, rest1) <- Word8 -> ByteString -> String -> Either String (Word8, ByteString)
takeExpectedByte Word8
0x08 ByteString
bs String
"expected tag 0x08 for field 1"
  _ <- pure tag1
  (typeVal, rest2) <- decodeUvarint rest1
  kt <- keyTypeFromProto typeVal
  -- Field 2: expect tag 0x12
  (_, rest3) <- takeExpectedByte 0x12 rest2 "expected tag 0x12 for field 2"
  (dataLen, rest4) <- decodeUvarint rest3
  let len = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
dataLen :: Int
  case compare (BS.length rest4) len of
    Ordering
LT -> String -> Either String (KeyType, ByteString)
forall a b. a -> Either a b
Left (String -> Either String (KeyType, ByteString))
-> String -> Either String (KeyType, ByteString)
forall a b. (a -> b) -> a -> b
$ String
ctx String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
": not enough bytes for key data"
    -- The peer-ids spec requires canonical serialization (Peer IDs are
    -- derived from these bytes); trailing data would give one key several
    -- byte representations, so reject it.
    Ordering
GT -> String -> Either String (KeyType, ByteString)
forall a b. a -> Either a b
Left (String -> Either String (KeyType, ByteString))
-> String -> Either String (KeyType, ByteString)
forall a b. (a -> b) -> a -> b
$ String
ctx String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
": trailing bytes after key data (non-canonical encoding)"
    Ordering
EQ -> (KeyType, ByteString) -> Either String (KeyType, ByteString)
forall a b. b -> Either a b
Right (KeyType
kt, ByteString
rest4)
  where
    takeExpectedByte :: Word8 -> ByteString -> String -> Either String (Word8, ByteString)
    takeExpectedByte :: Word8 -> ByteString -> String -> Either String (Word8, ByteString)
takeExpectedByte Word8
expected ByteString
input String
msg
      | ByteString -> Bool
BS.null ByteString
input = String -> Either String (Word8, ByteString)
forall a b. a -> Either a b
Left (String -> Either String (Word8, ByteString))
-> String -> Either String (Word8, ByteString)
forall a b. (a -> b) -> a -> b
$ String
ctx String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
": " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
msg String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" (empty input)"
      | HasCallStack => ByteString -> Word8
ByteString -> Word8
BS.head ByteString
input Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
expected =
          String -> Either String (Word8, ByteString)
forall a b. a -> Either a b
Left (String -> Either String (Word8, ByteString))
-> String -> Either String (Word8, ByteString)
forall a b. (a -> b) -> a -> b
$ String
ctx String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
": " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
msg String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" (got 0x" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Word8 -> String -> String
forall a. Integral a => a -> String -> String
showHex (HasCallStack => ByteString -> Word8
ByteString -> Word8
BS.head ByteString
input) String
")"
      | Bool
otherwise = (Word8, ByteString) -> Either String (Word8, ByteString)
forall a b. b -> Either a b
Right (Word8
expected, HasCallStack => ByteString -> ByteString
ByteString -> ByteString
BS.tail ByteString
input)