ml_kem/
crypto.rs

1#![allow(dead_code)]
2
3use crate::{B32, param::CbdSamplingSize};
4use module_lattice::EncodedPolynomial;
5use sha3::{Digest, Sha3_256, Sha3_512};
6use shake::{ExtendableOutput, Shake128, Shake256, Update, XofReader};
7
8pub(crate) fn G(inputs: &[impl AsRef<[u8]>]) -> (B32, B32) {
9    let mut h = Sha3_512::new();
10    for x in inputs {
11        Digest::update(&mut h, x);
12    }
13    let out = h.finalize();
14
15    let mut a = B32::default();
16    let mut b = B32::default();
17
18    a.copy_from_slice(&out[..32]);
19    b.copy_from_slice(&out[32..]);
20    (a, b)
21}
22
23pub(crate) fn H(x: impl AsRef<[u8]>) -> B32 {
24    let mut h = Sha3_256::new();
25    Digest::update(&mut h, x);
26
27    // This odd conversion is needed because the `sha3` crate links against an old version of
28    // the `generic-array` crate.  It should be pretty cheap though, since there's only one
29    // allocation / no copies.
30    let mut out = B32::default();
31    h.finalize_into(&mut out);
32    out
33}
34
35pub(crate) fn J(inputs: &[impl AsRef<[u8]>]) -> B32 {
36    let mut h = Shake256::default();
37    for x in inputs {
38        h.update(x.as_ref());
39    }
40    let mut r = h.finalize_xof();
41
42    let mut out = B32::default();
43    r.read(&mut out);
44    out
45}
46
47pub(crate) type PrfOutput<Eta> = EncodedPolynomial<<Eta as CbdSamplingSize>::SampleSize>;
48
49pub(crate) fn PRF<Eta>(s: &B32, b: u8) -> PrfOutput<Eta>
50where
51    Eta: CbdSamplingSize,
52{
53    let mut h = Shake256::default();
54    h.update(s.as_ref());
55    h.update(&[b]);
56    let mut r = h.finalize_xof();
57
58    let mut out = PrfOutput::<Eta>::default();
59    r.read(&mut out);
60    out
61}
62
63pub(crate) fn XOF(rho: &B32, i: u8, j: u8) -> impl XofReader {
64    let mut h = Shake128::default();
65    h.update(rho);
66    h.update(&[i, j]);
67    h.finalize_xof()
68}
69
70// // A Go script to generate the test vector outputs
71//
72// package main
73//
74// import (
75// 	"fmt"
76// 	"golang.org/x/crypto/sha3"
77// )
78//
79// func main() {
80// 	// G: B* -> B32 || B32 = SHA3_512(c)
81//   msgG := []byte("Input to an invocation of G")
82//   hG := sha3.New512()
83//   hG.Write(msgG)
84//   fmt.Printf("G: %x\n", hG.Sum(nil))
85//
86//   // H: B* -> B32 = SHA3_256(s)
87//   msgH := []byte("Input to an invocation of H")
88//   hH := sha3.New256()
89//   hH.Write(msgH)
90//   fmt.Printf("H: %x\n", hH.Sum(nil))
91//
92//   // J: B* -> B32 = SHAKE256(s, 32)
93//   msgJ := []byte("Input to an invocation of J")
94//   outJ := make([]byte, 32)
95//   sha3.ShakeSum256(outJ, msgJ)
96//   fmt.Printf("J: %x\n", outJ)
97//
98//   // PRF<2>: B32 x B -> B64eta = SHAKE256(s || b, 64 * eta)
99//   msgPRF2s := []byte("Input s to an invocation of PRF2")
100//   msgPRF2b := []byte("b")
101//   msgPRF2 := append(msgPRF2s, msgPRF2b...)
102//   outPRF2 := make([]byte, 64 * 2)
103//   sha3.ShakeSum256(outPRF2, msgPRF2)
104//   fmt.Printf("PRF<2>: %x\n", outPRF2)
105//
106//   // PRF<3>: B33 x B -> B64eta = SHAKE256(s || b, 64 * eta)
107//   msgPRF3s := []byte("Input s to an invocation of PRF3")
108//   msgPRF3b := []byte("b")
109//   msgPRF3 := append(msgPRF3s, msgPRF3b...)
110//   outPRF3 := make([]byte, 64 * 3)
111//   sha3.ShakeSum256(outPRF3, msgPRF3)
112//   fmt.Printf("PRF<3>: %x\n", outPRF3)
113//
114//   // XOF: B32 x B x B -> B* = SHAKE128(rho || i || j)
115//   msgXOFrho := []byte("Input rho, to an XOF invocation!")
116//   msgXOFi := []byte("i")
117//   msgXOFj := []byte("j")
118//   msgXOF := append(append(msgXOFrho, msgXOFi...), msgXOFj...)
119//   outXOF := make([]byte, 32)
120//   sha3.ShakeSum128(outXOF, msgXOF)
121//   fmt.Printf("XOF: %x\n", outXOF)
122//
123// }
124
125#[cfg(test)]
126mod test {
127    use super::*;
128    use array::typenum::{U2, U3};
129    use hex_literal::hex;
130
131    #[test]
132    fn g() {
133        let msg1 = "Input to ".as_bytes();
134        let msg2 = "an invocation of G".as_bytes();
135        let (actualA, actualB) = G(&[msg1, msg2]);
136        let expectedA = hex!("07dfced2a3a3feb3277cee1709818828ea6d2f42800152e9c312e848122231c2");
137        let expectedB = hex!("272969098a1bbd5a0a9844e2f89f206d8f7f4599e36aecaa4793af400fd880d8");
138        assert_eq!(actualA, expectedA);
139        assert_eq!(actualB, expectedB);
140    }
141
142    #[test]
143    fn h() {
144        let msg = "Input to an invocation of H".as_bytes();
145        let actual = H(msg);
146        let expected = hex!("0ee3ce94213d7dd0069b24b8b15cdd0bcf8eb1c6b3c21c441dc6a19e979cc7eb");
147        assert_eq!(actual, expected);
148    }
149
150    #[test]
151    fn j() {
152        let msg1 = "Input to ".as_bytes();
153        let msg2 = "an invocation of J".as_bytes();
154        let actual = J(&[msg1, msg2]);
155        let expected = hex!("a5292293d70c8eca049cbb475c48fabd625ed2b20785a18248504d3741196b52");
156        assert_eq!(actual, expected);
157    }
158
159    #[test]
160    fn prf() {
161        let s = B32::try_from("Input s to an invocation of PRF2".as_bytes())
162            .expect("Failed to create B32 from slice");
163        let b = b'b';
164        let actual = PRF::<U2>(&s, b);
165        let expected = hex!(
166            "54c002415c2219b564d5c17b0df0c82f83ddf3fdecc7d814ed5d85457c06c2c3\
167             ed0b0584f926dffb1e57c6105f8604e81c4605b93f8284e44585104101042075\
168             568113c861516d91bed227638654fc7f872df205c113b8364091755b62284eec\
169             a6124f2cd4c1cdf598cb8324a4f373470a8f81ee618c75cc33f66facee01c213"
170        );
171        assert_eq!(actual, expected);
172
173        let s = B32::try_from("Input s to an invocation of PRF3".as_bytes())
174            .expect("Failed to create B32 from slice");
175        let b = b'b';
176        let actual = PRF::<U3>(&s, b);
177        let expected = hex!(
178            "5e12028f67479b862a12713cda833e21b8ccd51bff9ddc2bfb9ab2910a9dc2e6\
179             c58264a3f51ccc9ef4ff936a15505e016f60c36ffe300be01b9fb12eacd57867\
180             0873c24709d6146b42c42a07873522eac100d61942ae53e73fbf9095b29b1ab7\
181             169e954213c062703dad88c1c5f57f92af143f0364fe057b134b54ea8a55d94c\
182             67764b3fc6b37376453978b8f0caeb6b18c188c28ee8681e28339477e042d5a1\
183             b4a12deb1de8b9dad026b4e323e03973ffbe25dd511eed5460d22a9851cfc220"
184        );
185        assert_eq!(actual, expected);
186    }
187
188    #[test]
189    fn xof() {
190        let rho = B32::try_from("Input rho, to an XOF invocation!".as_bytes())
191            .expect("Failed to create B32 from slice");
192        let i = b'i';
193        let j = b'j';
194
195        let mut reader = XOF(&rho, i, j);
196        let mut actual = [0u8; 32];
197        reader.read(&mut actual);
198
199        let expected = hex!("0d2c3e65f754d074cb366cf1b099ae105cc40f018342509f15f1ba8a1a4144cb");
200        assert_eq!(actual, expected);
201    }
202}
ml_kem/crypto.rs

ml_kem/
crypto.rs
