Argyle-Software / kyber

A rust implementation of the Kyber post-quantum KEM
https://docs.rs/pqc_kyber/
Apache License 2.0
163 stars 37 forks source link

gen_matrix is incorrect #77

Closed bwesterb closed 1 year ago

bwesterb commented 1 year ago

The current implementation of gen_matrix, given below, is incorrect. The problem is in the case that the initial BUFLEN=504 bytes squeezed from the XOF are not enough. Then only a single extra block is squeezed from the XOF, but the whole buffer is used for rejection sampling.

Compare your Rust implementation:

fn gen_matrix(a: &mut [Polyvec], seed: &[u8], transposed: bool)
{ 
  let mut ctr;
  // 530 is expected number of required bytes
  const GEN_MATRIX_NBLOCKS: usize = 
    (12*KYBER_N/8*(1 << 12)/KYBER_Q + XOF_BLOCKBYTES)/XOF_BLOCKBYTES;
  const BUFLEN: usize = GEN_MATRIX_NBLOCKS*XOF_BLOCKBYTES;
  let mut buf = [0u8; BUFLEN+2];
  let mut off: usize;
  let mut state = XofState::new();

  for i in 0..KYBER_K {
    for j in 0..KYBER_K {
      if transposed {
        xof_absorb(&mut state, seed, i as u8, j as u8);
      }
      else {
        xof_absorb(&mut state, seed, j as u8, i as u8);
      }
      xof_squeezeblocks(&mut buf, GEN_MATRIX_NBLOCKS, &mut state);
      ctr = rej_uniform(&mut a[i].vec[j].coeffs, KYBER_N, &buf, BUFLEN);

      while ctr < KYBER_N
      {
        off = BUFLEN % 3;
        for k in 0..off {
          buf[k] = buf[BUFLEN - off + k];
        }
        xof_squeezeblocks(&mut buf[off..], 1, &mut state);
        ctr += rej_uniform(&mut a[i].vec[j].coeffs[ctr..], KYBER_N - ctr, &buf, BUFLEN);
      }
    }
  }
}

To the reference implementation where buflen is correctly adjusted:

void gen_matrix(polyvec *a, const uint8_t seed[KYBER_SYMBYTES], int transposed)
{
  unsigned int ctr, i, j, k;
  unsigned int buflen, off;
  uint8_t buf[GEN_MATRIX_NBLOCKS*XOF_BLOCKBYTES+2];
  xof_state state;

  for(i=0;i<KYBER_K;i++) {
    for(j=0;j<KYBER_K;j++) {
      if(transposed)
        xof_absorb(&state, seed, i, j);
      else
        xof_absorb(&state, seed, j, i);

      xof_squeezeblocks(buf, GEN_MATRIX_NBLOCKS, &state);
      buflen = GEN_MATRIX_NBLOCKS*XOF_BLOCKBYTES;
      ctr = rej_uniform(a[i].vec[j].coeffs, KYBER_N, buf, buflen);

      while(ctr < KYBER_N) {
        off = buflen % 3;
        for(k = 0; k < off; k++)
          buf[k] = buf[buflen - off + k];
        xof_squeezeblocks(buf + off, 1, &state);
        buflen = off + XOF_BLOCKBYTES;
        ctr += rej_uniform(a[i].vec[j].coeffs + ctr, KYBER_N - ctr, buf, buflen);
      }
    }
  }
}

The issue for regular Kyber occurs with probability approximately 2^-105.

bwesterb commented 1 year ago

For 90s Kyber (which no one should use) you hit it with probability 2^-38.

mberry commented 1 year ago

Thanks for catching this.