Rewrite NIKEs and KEMs to use structure values

katzenpost · Aug 31, 2024 · 692fcd1 · 692fcd1
1 parent 005f185
commit 692fcd1
Show file tree

Hide file tree

Showing 11 changed files with 280 additions and 202 deletions.
diff --git a/CryptWalker.lean b/CryptWalker.lean
@@ -2,6 +2,7 @@ import CryptWalker.nike.x25519
 import CryptWalker.nike.x448
 import CryptWalker.nike.x41417
 import CryptWalker.nike.nike
+import CryptWalker.nike.schemes
 import CryptWalker.hash.Sha2
 import CryptWalker.hash.Sha512
 import CryptWalker.util.newnat

diff --git a/CryptWalker/kem/adapter.lean b/CryptWalker/kem/adapter.lean
@@ -22,68 +22,71 @@ structure PublicKey where
 
 instance : kem.Key PrivateKey where
   encode : PrivateKey → ByteArray := fun (key : PrivateKey) => key.data
-  decode : ByteArray → Option PrivateKey := fun (bytes : ByteArray) => some (PrivateKey.mk bytes)
+  decode (_kem : KEM) (bytes : ByteArray) : Option PrivateKey :=
+    some (PrivateKey.mk bytes)
 
 instance : kem.Key PublicKey where
   encode : PublicKey → ByteArray := fun (key : PublicKey) => key.data
-  decode : ByteArray → Option PublicKey := fun (bytes : ByteArray) => some (PublicKey.mk bytes)
+  decode (_kem : KEM) (bytes : ByteArray) : Option PublicKey :=
+    some (PublicKey.mk bytes)
 
 
-class Adapter (α : Type) [NIKE α] where
+structure Adapter where
   hash : ByteArray → ByteArray
+  nike : NIKE
 
-instance {α : Type} [nikeInstance : NIKE α] [adapter : Adapter α] : KEM α where
+instance (adapter : Adapter) : KEM where
   PublicKeyType := PublicKey
   PrivateKeyType := PrivateKey
 
-  name : String := NIKE.name α
+  name : String := adapter.nike.name
 
   generateKeyPair : IO (PublicKey × PrivateKey) := do
-    let keyPair ← nikeInstance.generateKeyPair
+    let keyPair ← adapter.nike.generateKeyPair
     let pubkey := keyPair.1
     let privkey := keyPair.2
-    let pubkeyData := NIKE.encodePublicKey pubkey
-    let privkeyData := NIKE.encodePrivateKey privkey
+    let pubkeyData := adapter.nike.encodePublicKey pubkey
+    let privkeyData := adapter.nike.encodePrivateKey privkey
     pure (PublicKey.mk pubkeyData, PrivateKey.mk privkeyData)
 
-  encapsulate (_self : α) (theirPubKey : PublicKey) : IO (ByteArray × ByteArray) := do
-    let (pubkey, privkey) ← NIKE.generateKeyPair
-    let theirNikePubKeyOpt : Option (NIKE.PublicKeyType α) := NIKE.decodePublicKey theirPubKey.data
+  encapsulate (theirPubKey : PublicKey) : IO (ByteArray × ByteArray) := do
+    let (pubkey, privkey) ← adapter.nike.generateKeyPair
+    let theirNikePubKeyOpt : Option (adapter.nike.PublicKeyType) := adapter.nike.decodePublicKey theirPubKey.data
     match theirNikePubKeyOpt with
     | none => panic! "Failed to decode NIKE public key"
     | some theirNikePubKey =>
-      let ss1 := NIKE.groupAction privkey theirNikePubKey
-      let ss1Bytes := NIKE.encodePublicKey ss1
-      let pubkeyBytes := NIKE.encodePublicKey pubkey
+      let ss1 := adapter.nike.groupAction privkey theirNikePubKey
+      let ss1Bytes := adapter.nike.encodePublicKey ss1
+      let pubkeyBytes := adapter.nike.encodePublicKey pubkey
       let blob := ByteArray.append ss1Bytes $ ByteArray.append theirPubKey.data pubkeyBytes
       let ss2 := adapter.hash blob
-      let ciphertext := NIKE.encodePublicKey pubkey
+      let ciphertext := adapter.nike.encodePublicKey pubkey
       pure (ciphertext, ss2)
 
-  decapsulate (_self : α) (privKey : PrivateKey) (ct : ByteArray) : ByteArray :=
-    let theirPubKeyOpt : Option (NIKE.PublicKeyType α) := NIKE.decodePublicKey ct
+  decapsulate (privKey : PrivateKey) (ct : ByteArray) : ByteArray :=
+    let theirPubKeyOpt : Option (adapter.nike.PublicKeyType) := adapter.nike.decodePublicKey ct
     match theirPubKeyOpt with
     | none => panic! "adapter decap failure: failed to decode NIKE public key"
     | some theirPubKey =>
-      let myPrivKeyOpt : Option (NIKE.PrivateKeyType α) := NIKE.decodePrivateKey privKey.data
+      let myPrivKeyOpt : Option (adapter.nike.PrivateKeyType) := adapter.nike.decodePrivateKey privKey.data
       match myPrivKeyOpt with
       | none => panic! "adapter decap failure: failed to decode NIKE private key"
       | some myPrivKey =>
-        let ss1 := NIKE.groupAction myPrivKey theirPubKey
-        let a := NIKE.encodePublicKey ss1
-        let myPubKey := NIKE.derivePublicKey myPrivKey
-        let myPubKeyBytes := NIKE.encodePublicKey myPubKey
-        let theirPubKeyBytes := NIKE.encodePublicKey theirPubKey
+        let ss1 := adapter.nike.groupAction myPrivKey theirPubKey
+        let a := adapter.nike.encodePublicKey ss1
+        let myPubKey := adapter.nike.derivePublicKey myPrivKey
+        let myPubKeyBytes := adapter.nike.encodePublicKey myPubKey
+        let theirPubKeyBytes := adapter.nike.encodePublicKey theirPubKey
         let b := ByteArray.append myPubKeyBytes theirPubKeyBytes
         let blob := ByteArray.append a b
         adapter.hash blob
 
-  privateKeySize := nikeInstance.privateKeySize
-  publicKeySize := nikeInstance.publicKeySize
+  privateKeySize := adapter.nike.privateKeySize
+  publicKeySize := adapter.nike.publicKeySize
 
-  encodePrivateKey := fun (sk : PrivateKey) => kem.Key.encode sk
-  decodePrivateKey := fun (bytes : ByteArray) => kem.Key.decode bytes
-  encodePublicKey := fun (pk : PublicKey) => kem.Key.encode pk
-  decodePublicKey := fun (bytes : ByteArray) => kem.Key.decode bytes
+  encodePrivateKey (sk : PrivateKey) : ByteArray := sk.data
+  decodePrivateKey (bytes : ByteArray) : Option PrivateKey := some {data := bytes}
+  encodePublicKey (pk : PublicKey) : ByteArray := pk.data
+  decodePublicKey (bytes : ByteArray) : Option PublicKey := some {data := bytes}
 
 end CryptWalker.kem.adapter
diff --git a/CryptWalker/kem/combiner.lean b/CryptWalker/kem/combiner.lean
@@ -5,23 +5,71 @@ SPDX-License-Identifier: AGPL-3.0-only
 
 -- Security preserving KEM combiner
 
+import Mathlib.Data.ByteArray
+
 import CryptWalker.kem.kem
 namespace CryptWalker.kem.combiner
 open CryptWalker.kem.kem
 
+-- alternative combiner would use a type parameter
+
+structure Combiner where
+  hash : ByteArray → ByteArray
+  KEMs : List (Σ α : Type, KEM α × α)
+
+
 structure PrivateKey where
-  data : ByteArray
+  data : List ByteArray
 
 structure PublicKey where
-  data : ByteArray
+  data : List ByteArray
+
+def splitByteArray (bytes : ByteArray) (n : Nat) : ByteArray × ByteArray :=
+  let part1 := bytes.extract 0 n
+  let part2 := bytes.extract n bytes.size
+  (part1, part2)
+
+def splitByteArrayIntoChunks (bytes : ByteArray) (sizes : List Nat) : Option (List ByteArray) :=
+  let rec aux (bytes : ByteArray) (sizes : List Nat) (acc : List ByteArray) : Option (List ByteArray) :=
+    match sizes with
+    | [] =>
+      if bytes.isEmpty then
+        some acc.reverse
+      else
+        none
+    | size :: sizesTail =>
+      if bytes.size < size then
+        none
+      else
+        let (part1, part2) := splitByteArray bytes size
+        aux part2 sizesTail (part1 :: acc)
+  aux bytes sizes []
 
 instance : kem.Key PrivateKey where
-  encode : PrivateKey → ByteArray := fun (key : PrivateKey) => key.data
-  decode : ByteArray → Option PrivateKey := fun (bytes : ByteArray) => some (PrivateKey.mk bytes)
+  encode : PrivateKey → ByteArray := fun (key : PrivateKey) =>
+    key.data.foldl (init:= ByteArray.empty) fun acc byteArray => acc ++ byteArray
+
+  decode {α : Type} (kem : KEM α Combiner) (c : Combiner) (bytes : ByteArray) : Option PrivateKey :=
+    if bytes.size ≠ kem.privateKeySize then
+      none
+    else
+      let sizes := c.KEMs.foldl (f := fun acc x => acc ++ [x.snd.fst.privateKeySize]) []
+      match splitByteArrayIntoChunks bytes sizes with
+      | none => none
+      | some keys => some { data := keys}
+
 
 instance : kem.Key PublicKey where
-  encode : PublicKey → ByteArray := fun (key : PublicKey) => key.data
-  decode : ByteArray → Option PublicKey := fun (bytes : ByteArray) => some (PublicKey.mk bytes)
+  encode : PublicKey → ByteArray := fun (key : PublicKey) =>
+    key.data.foldl (init:= ByteArray.empty) fun acc byteArray => acc ++ byteArray
+  decode : ByteArray → Option PublicKey := fun (bytes : ByteArray) =>
+    if bytes.size ≠ kem.publicKeySize then
+      none
+    blah
+
+    --some (PublicKey.mk bytes)
+    -- FIXME
+    none
 
 
 

diff --git a/CryptWalker/kem/kem.lean b/CryptWalker/kem/kem.lean
@@ -1,31 +1,29 @@
-/-
-SPDX-FileCopyrightText: Copyright (C) 2024 David Stainton
-SPDX-License-Identifier: AGPL-3.0-only
- -/
+import Batteries.Classes.SatisfiesM
 
 namespace CryptWalker.kem.kem
 
-class Key (key : Type) where
-  encode : key → ByteArray
-  decode : ByteArray → Option key
+  structure KEM where
+    PublicKeyType : Type
+    PrivateKeyType : Type
 
-class PrivateKey (privkey : Type) extends Key privkey
+    name : String
+    privateKeySize : Nat
+    publicKeySize : Nat
 
-class PublicKey (pubkey : Type) extends Key pubkey
+    generateKeyPair : IO (PublicKeyType × PrivateKeyType)
+    encapsulate : PublicKeyType → IO (ByteArray × ByteArray)
+    decapsulate : PrivateKeyType → ByteArray → ByteArray
+    encodePrivateKey : PrivateKeyType → ByteArray
+    decodePrivateKey : ByteArray → Option PrivateKeyType
+    encodePublicKey : PublicKeyType → ByteArray
+    decodePublicKey : ByteArray → Option PublicKeyType
 
-class KEM (scheme : Type) where
-  PublicKeyType : Type
-  PrivateKeyType : Type
+  structure Key (key : Type) where
+    encode : key → ByteArray
+    decode : KEM → ByteArray → Option key
 
-  name : String
-  generateKeyPair : IO (PublicKeyType × PrivateKeyType)
-  encapsulate : scheme → PublicKeyType → IO (ByteArray × ByteArray)
-  decapsulate : scheme → PrivateKeyType → ByteArray → ByteArray
-  privateKeySize : Nat
-  publicKeySize : Nat
-  encodePrivateKey : PrivateKeyType → ByteArray
-  decodePrivateKey : ByteArray → Option PrivateKeyType
-  encodePublicKey : PublicKeyType → ByteArray
-  decodePublicKey : ByteArray → Option PublicKeyType
+  structure PrivateKey (privkey : Type) extends Key privkey
+
+  structure PublicKey (pubkey : Type) extends Key pubkey
 
 end CryptWalker.kem.kem
diff --git a/CryptWalker/kem/schemes.lean b/CryptWalker/kem/schemes.lean
@@ -7,40 +7,55 @@ import CryptWalker.kem.kem
 import CryptWalker.kem.adapter
 import CryptWalker.hash.Sha2
 
-open CryptWalker.nike.x25519
-open CryptWalker.nike.x448
-open CryptWalker.nike.x41417
+open CryptWalker.nike
 open CryptWalker.nike.nike
 open CryptWalker.kem.kem
 open CryptWalker.kem.adapter
 open CryptWalker.hash.Sha2
 
 namespace CryptWalker.kem.schemes
 
-def defaultHash := fun x => Sha256.hash x
-
-instance : Adapter X25519Scheme where
-  hash := defaultHash
-
-instance : Adapter X448Scheme where
-  hash := defaultHash
-
-instance : Adapter X41417Scheme where
-  hash := defaultHash
-
-def X25519AsKEM : KEM X25519Scheme := inferInstance
-def X448AsKEM : KEM X448Scheme := inferInstance
-def X41417AsKEM : KEM X41417Scheme := inferInstance
-
-def X25519Instance : X25519Scheme := {}
-def X448Instance : X448Scheme := {}
-def X41417Instance : X41417Scheme := {}
-
-def Schemes : List (Σ α : Type, KEM α × α) :=
-  [
-    ⟨X25519Scheme, (X25519AsKEM, X25519Instance)⟩,
-    ⟨X448Scheme, (X448AsKEM, X448Instance)⟩,
-    ⟨X41417Scheme, (X41417AsKEM, X41417Instance)⟩,
-  ]
+def toKEM (adapter : Adapter) : KEM :=
+  {
+    PublicKeyType := PublicKey,
+    PrivateKeyType := PrivateKey,
+    name := adapter.nike.name,
+    generateKeyPair := do
+      let keyPair ← adapter.nike.generateKeyPair
+      let pubkey := PublicKey.mk (adapter.nike.encodePublicKey keyPair.1)
+      let privkey := PrivateKey.mk (adapter.nike.encodePrivateKey keyPair.2)
+      pure (pubkey, privkey),
+    encapsulate := fun theirPubKey => do
+      let (pubkey, privkey) ← adapter.nike.generateKeyPair
+      match adapter.nike.decodePublicKey theirPubKey.data with
+      | none => panic! "type coercion failure"
+      | some pubkey2 =>
+        let ss1 := adapter.nike.groupAction privkey pubkey2
+        let ss2 := adapter.hash (adapter.nike.encodePublicKey ss1)
+        let ciphertext := adapter.nike.encodePublicKey pubkey
+        pure (ciphertext, ss2),
+    decapsulate := fun privKey ct =>
+      match adapter.nike.decodePublicKey ct with
+      | none => panic! "type coercion failure"
+      | some pubkey2 =>
+        match adapter.nike.decodePrivateKey privKey.data with
+        | none => panic! "type coercion failure"
+        | some privkey2 =>
+          let ss1 := adapter.nike.groupAction privkey2 pubkey2
+          adapter.hash (adapter.nike.encodePublicKey ss1),
+    privateKeySize := adapter.nike.privateKeySize,
+    publicKeySize := adapter.nike.publicKeySize,
+    encodePrivateKey := fun sk => sk.data,
+    decodePrivateKey := fun bytes => some {data := bytes},
+    encodePublicKey := fun pk => pk.data,
+    decodePublicKey := fun bytes => some {data := bytes}
+  }
+
+def Schemes : List KEM :=
+[
+    toKEM $ Adapter.mk Sha256.hash x25519.Scheme,
+    toKEM $ Adapter.mk Sha256.hash x448.Scheme,
+    toKEM $ Adapter.mk Sha256.hash x41417.Scheme
+]
 
 end CryptWalker.kem.schemes
diff --git a/CryptWalker/nike/nike.lean b/CryptWalker/nike/nike.lean
@@ -2,18 +2,11 @@
 SPDX-FileCopyrightText: Copyright (C) 2024 David Stainton
 SPDX-License-Identifier: AGPL-3.0-only
  -/
+import Batteries.Classes.SatisfiesM
 
 namespace CryptWalker.nike.nike
 
-class Key (key : Type) where
-  encode : key → ByteArray
-  decode : ByteArray → Option key
-
-class PrivateKey (privkey : Type) extends Key privkey
-
-class PublicKey (pubkey : Type) extends Key pubkey
-
-class NIKE (scheme : Type) where
+structure NIKE where
   PublicKeyType : Type
   PrivateKeyType : Type
 
@@ -30,4 +23,12 @@ class NIKE (scheme : Type) where
   encodePublicKey : PublicKeyType → ByteArray
   decodePublicKey : ByteArray → Option PublicKeyType
 
+structure Key (key : Type) where
+  encode : key → ByteArray
+  decode : NIKE → ByteArray → Option key
+
+structure PrivateKey (privkey : Type) extends Key privkey
+
+structure PublicKey (pubkey : Type) extends Key pubkey
+
 end CryptWalker.nike.nike
diff --git a/CryptWalker/nike/schemes.lean b/CryptWalker/nike/schemes.lean
@@ -0,0 +1,15 @@
+import CryptWalker.nike.x25519
+import CryptWalker.nike.x448
+import CryptWalker.nike.x41417
+import CryptWalker.nike.nike
+
+namespace CryptWalker.nike.schemes
+
+def Schemes : List CryptWalker.nike.nike.NIKE :=
+  [
+    CryptWalker.nike.x25519.Scheme,
+    CryptWalker.nike.x448.Scheme,
+    CryptWalker.nike.x41417.Scheme,
+  ]
+
+end CryptWalker.nike.schemes