Browse Source

Initial import.

Yawning Angel 1 year ago
commit
47a9d7cb34
19 changed files with 2097 additions and 0 deletions
  1. 2 0
      .gitignore
  2. 122 0
      LICENSE
  3. 96 0
      cbd.go
  4. 22 0
      doc.go
  5. 279 0
      indcpa.go
  6. 183 0
      kem.go
  7. 185 0
      kem_test.go
  8. 294 0
      kem_vectors_test.go
  9. 241 0
      kex.go
  10. 83 0
      kex_test.go
  11. 65 0
      ntt.go
  12. 116 0
      params.go
  13. 146 0
      poly.go
  14. 110 0
      polyvec.go
  15. 87 0
      precomp.go
  16. 43 0
      reduce.go
  17. 1 0
      testdata/.gitignore
  18. 17 0
      testdata/README.testdata
  19. 5 0
      testdata/compactVectors.json

+ 2 - 0
.gitignore

@@ -0,0 +1,2 @@
+*.swp
+*~

+ 122 - 0
LICENSE

@@ -0,0 +1,122 @@
+Creative Commons Legal Code
+
+CC0 1.0 Universal
+
+    CREATIVE COMMONS CORPORATION IS NOT A LAW FIRM AND DOES NOT PROVIDE
+    LEGAL SERVICES. DISTRIBUTION OF THIS DOCUMENT DOES NOT CREATE AN
+    ATTORNEY-CLIENT RELATIONSHIP. CREATIVE COMMONS PROVIDES THIS
+    INFORMATION ON AN "AS-IS" BASIS. CREATIVE COMMONS MAKES NO WARRANTIES
+    REGARDING THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS
+    PROVIDED HEREUNDER, AND DISCLAIMS LIABILITY FOR DAMAGES RESULTING FROM
+    THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS PROVIDED
+    HEREUNDER.
+
+Statement of Purpose
+
+The laws of most jurisdictions throughout the world automatically confer
+exclusive Copyright and Related Rights (defined below) upon the creator
+and subsequent owner(s) (each and all, an "owner") of an original work of
+authorship and/or a database (each, a "Work").
+
+Certain owners wish to permanently relinquish those rights to a Work for
+the purpose of contributing to a commons of creative, cultural and
+scientific works ("Commons") that the public can reliably and without fear
+of later claims of infringement build upon, modify, incorporate in other
+works, reuse and redistribute as freely as possible in any form whatsoever
+and for any purposes, including without limitation commercial purposes.
+These owners may contribute to the Commons to promote the ideal of a free
+culture and the further production of creative, cultural and scientific
+works, or to gain reputation or greater distribution for their Work in
+part through the use and efforts of others.
+
+For these and/or other purposes and motivations, and without any
+expectation of additional consideration or compensation, the person
+associating CC0 with a Work (the "Affirmer"), to the extent that he or she
+is an owner of Copyright and Related Rights in the Work, voluntarily
+elects to apply CC0 to the Work and publicly distribute the Work under its
+terms, with knowledge of his or her Copyright and Related Rights in the
+Work and the meaning and intended legal effect of CC0 on those rights.
+
+1. Copyright and Related Rights. A Work made available under CC0 may be
+protected by copyright and related or neighboring rights ("Copyright and
+Related Rights"). Copyright and Related Rights include, but are not
+limited to, the following:
+
+  i. the right to reproduce, adapt, distribute, perform, display,
+     communicate, and translate a Work;
+ ii. moral rights retained by the original author(s) and/or performer(s);
+iii. publicity and privacy rights pertaining to a person's image or
+     likeness depicted in a Work;
+ iv. rights protecting against unfair competition in regards to a Work,
+     subject to the limitations in paragraph 4(a), below;
+  v. rights protecting the extraction, dissemination, use and reuse of data
+     in a Work;
+ vi. database rights (such as those arising under Directive 96/9/EC of the
+     European Parliament and of the Council of 11 March 1996 on the legal
+     protection of databases, and under any national implementation
+     thereof, including any amended or successor version of such
+     directive); and
+vii. other similar, equivalent or corresponding rights throughout the
+     world based on applicable law or treaty, and any national
+     implementations thereof.
+
+2. Waiver. To the greatest extent permitted by, but not in contravention
+of, applicable law, Affirmer hereby overtly, fully, permanently,
+irrevocably and unconditionally waives, abandons, and surrenders all of
+Affirmer's Copyright and Related Rights and associated claims and causes
+of action, whether now known or unknown (including existing as well as
+future claims and causes of action), in the Work (i) in all territories
+worldwide, (ii) for the maximum duration provided by applicable law or
+treaty (including future time extensions), (iii) in any current or future
+medium and for any number of copies, and (iv) for any purpose whatsoever,
+including without limitation commercial, advertising or promotional
+purposes (the "Waiver"). Affirmer makes the Waiver for the benefit of each
+member of the public at large and to the detriment of Affirmer's heirs and
+successors, fully intending that such Waiver shall not be subject to
+revocation, rescission, cancellation, termination, or any other legal or
+equitable action to disrupt the quiet enjoyment of the Work by the public
+as contemplated by Affirmer's express Statement of Purpose.
+
+3. Public License Fallback. Should any part of the Waiver for any reason
+be judged legally invalid or ineffective under applicable law, then the
+Waiver shall be preserved to the maximum extent permitted taking into
+account Affirmer's express Statement of Purpose. In addition, to the
+extent the Waiver is so judged Affirmer hereby grants to each affected
+person a royalty-free, non transferable, non sublicensable, non exclusive,
+irrevocable and unconditional license to exercise Affirmer's Copyright and
+Related Rights in the Work (i) in all territories worldwide, (ii) for the
+maximum duration provided by applicable law or treaty (including future
+time extensions), (iii) in any current or future medium and for any number
+of copies, and (iv) for any purpose whatsoever, including without
+limitation commercial, advertising or promotional purposes (the
+"License"). The License shall be deemed effective as of the date CC0 was
+applied by Affirmer to the Work. Should any part of the License for any
+reason be judged legally invalid or ineffective under applicable law, such
+partial invalidity or ineffectiveness shall not invalidate the remainder
+of the License, and in such case Affirmer hereby affirms that he or she
+will not (i) exercise any of his or her remaining Copyright and Related
+Rights in the Work or (ii) assert any associated claims and causes of
+action with respect to the Work, in either case contrary to Affirmer's
+express Statement of Purpose.
+
+4. Limitations and Disclaimers.
+
+ a. No trademark or patent rights held by Affirmer are waived, abandoned,
+    surrendered, licensed or otherwise affected by this document.
+ b. Affirmer offers the Work as-is and makes no representations or
+    warranties of any kind concerning the Work, express, implied,
+    statutory or otherwise, including without limitation warranties of
+    title, merchantability, fitness for a particular purpose, non
+    infringement, or the absence of latent or other defects, accuracy, or
+    the present or absence of errors, whether or not discoverable, all to
+    the greatest extent permissible under applicable law.
+ c. Affirmer disclaims responsibility for clearing rights of other persons
+    that may apply to the Work or any use thereof, including without
+    limitation any person's Copyright and Related Rights in the Work.
+    Further, Affirmer disclaims responsibility for obtaining any necessary
+    consents, permissions or other rights required for any use of the
+    Work.
+ d. Affirmer understands and acknowledges that Creative Commons is not a
+    party to this document and has no duty or obligation with respect to
+    this CC0 or use of the Work.
+

+ 96 - 0
cbd.go

@@ -0,0 +1,96 @@
+// cbd.go - Centered binomial distribution.
+//
+// To the extent possible under law, Yawning Angel has waived all copyright
+// and related or neighboring rights to the software, using the Creative
+// Commons "CC0" public domain dedication. See LICENSE or
+// <http://creativecommons.org/publicdomain/zero/1.0/> for full details.
+
+package kyber
+
+// Load bytes into a 64-bit integer in little-endian order.
+func loadLittleEndian(x []byte, bytes int) uint64 {
+	var r uint64
+	for i, v := range x[:bytes] {
+		r |= uint64(v) << (8 * uint(i))
+	}
+	return r
+}
+
+// Given an array of uniformly random bytes, compute polynomial with
+// coefficients distributed according to a centered binomial distribution
+// with parameter eta.
+func (p *poly) cbd(buf []byte, eta int) {
+	switch eta {
+	case 3:
+		var a, b [4]uint32
+		for i := 0; i < kyberN/4; i++ {
+			t := loadLittleEndian(buf[3*i:], 3)
+			var d uint32
+			for j := 0; j < 3; j++ {
+				d += uint32((t >> uint(j)) & 0x249249)
+			}
+
+			a[0] = d & 0x7
+			b[0] = (d >> 3) & 0x7
+			a[1] = (d >> 6) & 0x7
+			b[1] = (d >> 9) & 0x7
+			a[2] = (d >> 12) & 0x7
+			b[2] = (d >> 15) & 0x7
+			a[3] = (d >> 18) & 0x7
+			b[3] = (d >> 21)
+
+			p.coeffs[4*i+0] = uint16(a[0] + kyberQ - b[0])
+			p.coeffs[4*i+1] = uint16(a[1] + kyberQ - b[1])
+			p.coeffs[4*i+2] = uint16(a[2] + kyberQ - b[2])
+			p.coeffs[4*i+3] = uint16(a[3] + kyberQ - b[3])
+		}
+	case 4:
+		var a, b [4]uint32
+		for i := 0; i < kyberN/4; i++ {
+			t := loadLittleEndian(buf[4*i:], 4)
+			var d uint32
+			for j := 0; j < 4; j++ {
+				d += uint32((t >> uint(j)) & 0x11111111)
+			}
+
+			a[0] = d & 0xf
+			b[0] = (d >> 4) & 0xf
+			a[1] = (d >> 8) & 0xf
+			b[1] = (d >> 12) & 0xf
+			a[2] = (d >> 16) & 0xf
+			b[2] = (d >> 20) & 0xf
+			a[3] = (d >> 24) & 0xf
+			b[3] = (d >> 28)
+
+			p.coeffs[4*i+0] = uint16(a[0] + kyberQ - b[0])
+			p.coeffs[4*i+1] = uint16(a[1] + kyberQ - b[1])
+			p.coeffs[4*i+2] = uint16(a[2] + kyberQ - b[2])
+			p.coeffs[4*i+3] = uint16(a[3] + kyberQ - b[3])
+		}
+	case 5:
+		var a, b [4]uint64
+		for i := 0; i < kyberN/4; i++ {
+			t := loadLittleEndian(buf[5*i:], 5)
+			var d uint64
+			for j := 0; j < 5; j++ {
+				d += (t >> uint(j)) & 0x0842108421
+			}
+
+			a[0] = d & 0x1f
+			b[0] = (d >> 5) & 0x1f
+			a[1] = (d >> 10) & 0x1f
+			b[1] = (d >> 15) & 0x1f
+			a[2] = (d >> 20) & 0x1f
+			b[2] = (d >> 25) & 0x1f
+			a[3] = (d >> 30) & 0x1f
+			b[3] = (d >> 35)
+
+			p.coeffs[4*i+0] = uint16(a[0] + kyberQ - b[0])
+			p.coeffs[4*i+1] = uint16(a[1] + kyberQ - b[1])
+			p.coeffs[4*i+2] = uint16(a[2] + kyberQ - b[2])
+			p.coeffs[4*i+3] = uint16(a[3] + kyberQ - b[3])
+		}
+	default:
+		panic("kyber: eta must be in {3,4,5}")
+	}
+}

+ 22 - 0
doc.go

@@ -0,0 +1,22 @@
+// doc.go - Kyber godoc extras.
+//
+// To the extent possible under law, Yawning Angel has waived all copyright
+// and related or neighboring rights to the software, using the Creative
+// Commons "CC0" public domain dedication. See LICENSE or
+// <http://creativecommons.org/publicdomain/zero/1.0/> for full details.
+
+// Package kyber implements the Kyber IND-CCA2-secure key encapsulation
+// mechanism (KEM), based on the hardness of solving the learning-with-errors
+// (LWE) problem over module lattices as submitted to the NIST Post-Quantum
+// Cryptography project.
+//
+// This implementation is a port of the Public Domain reference implementation
+// by Joppe Bos, Léo Ducas, Eike Kiltz , Tancrède Lepoint, Vadim Lyubashevsky,
+// John Schanck, Peter Schwabe, Gregor Seiler, and Damien Stehlé.
+//
+// Additionally implementations of Kyber.AKE and Kyber.UAKE as presented in
+// the Kyber paper are included for users that seek an authenticated key
+// exchange.
+//
+// For more information, see https://pq-crystals.org/kyber/index.shtml.
+package kyber

+ 279 - 0
indcpa.go

@@ -0,0 +1,279 @@
+// indcpa.go - Kyber IND-CPA encryption.
+//
+// To the extent possible under law, Yawning Angel has waived all copyright
+// and related or neighboring rights to the software, using the Creative
+// Commons "CC0" public domain dedication. See LICENSE or
+// <http://creativecommons.org/publicdomain/zero/1.0/> for full details.
+
+package kyber
+
+import (
+	"io"
+
+	"golang.org/x/crypto/sha3"
+)
+
+// Serialize the public key as concatenation of the compressed and serialized
+// vector of polynomials pk and the public seed used to generate the matrix A.
+func packPublicKey(r []byte, pk *polyVec, seed []byte) {
+	pk.compress(r)
+	copy(r[pk.compressedSize():], seed[:SymSize])
+}
+
+// De-serialize and decompress public key from a byte array; approximate
+// inverse of packPublicKey.
+func unpackPublicKey(pk *polyVec, seed, packedPk []byte) {
+	pk.decompress(packedPk)
+
+	off := pk.compressedSize()
+	copy(seed, packedPk[off:off+SymSize])
+}
+
+// Serialize the ciphertext as concatenation of the compressed and serialized
+// vector of polynomials b and the compressed and serialized polynomial v.
+func packCiphertext(r []byte, b *polyVec, v *poly) {
+	b.compress(r)
+	v.compress(r[b.compressedSize():])
+}
+
+// De-serialize and decompress ciphertext from a byte array; approximate
+// inverse of packCiphertext.
+func unpackCiphertext(b *polyVec, v *poly, c []byte) {
+	b.decompress(c)
+	v.decompress(c[b.compressedSize():])
+}
+
+// Serialize the secret key.
+func packSecretKey(r []byte, sk *polyVec) {
+	sk.toBytes(r)
+}
+
+// De-serialize the secret key; inverse of packSecretKey.
+func unpackSecretKey(sk *polyVec, packedSk []byte) {
+	sk.fromBytes(packedSk)
+}
+
+// Deterministically generate matrix A (or the transpose of A) from a seed.
+// Entries of the matrix are polynomials that look uniformly random. Performs
+// rejection sampling on output of SHAKE-128.
+func genMatrix(a []polyVec, seed []byte, transposed bool) {
+	const (
+		shake128Rate = 168 // xof.BlockSize() is not a constant.
+		maxBlocks    = 4
+	)
+	var buf [shake128Rate * maxBlocks]byte
+
+	var extSeed [SymSize + 2]byte
+	copy(extSeed[:SymSize], seed)
+
+	xof := sha3.NewShake128()
+
+	for i, v := range a {
+		for j, p := range v.vec {
+			if transposed {
+				extSeed[SymSize] = byte(i)
+				extSeed[SymSize+1] = byte(j)
+			} else {
+				extSeed[SymSize] = byte(j)
+				extSeed[SymSize+1] = byte(i)
+			}
+
+			xof.Write(extSeed[:])
+			xof.Read(buf[:])
+
+			for ctr, pos, maxPos := 0, 0, len(buf); ctr < kyberN; {
+				val := (uint16(buf[pos]) | (uint16(buf[pos+1]) << 8)) & 0x1fff
+				if val < kyberQ {
+					p.coeffs[ctr] = val
+					ctr++
+				}
+				if pos += 2; pos == maxPos {
+					// On the unlikely chance 4 blocks is insufficient,
+					// incrementally squeeze out 1 block at a time.
+					xof.Read(buf[:shake128Rate])
+					pos, maxPos = 0, shake128Rate
+				}
+			}
+
+			xof.Reset()
+		}
+	}
+}
+
+type indcpaPublicKey struct {
+	packed []byte
+	h      [32]byte
+}
+
+func (pk *indcpaPublicKey) toBytes() []byte {
+	return pk.packed
+}
+
+func (pk *indcpaPublicKey) fromBytes(p *ParameterSet, b []byte) error {
+	if len(b) != p.indcpaPublicKeySize {
+		return ErrInvalidKeySize
+	}
+
+	pk.packed = make([]byte, len(b))
+	copy(pk.packed, b)
+	pk.h = sha3.Sum256(b)
+
+	return nil
+}
+
+type indcpaSecretKey struct {
+	packed []byte
+}
+
+func (sk *indcpaSecretKey) fromBytes(p *ParameterSet, b []byte) error {
+	if len(b) != p.indcpaSecretKeySize {
+		return ErrInvalidKeySize
+	}
+
+	sk.packed = make([]byte, len(b))
+	copy(sk.packed, b)
+
+	return nil
+}
+
+// Generates public and private key for the CPA-secure public-key encryption
+// scheme underlying Kyber.
+func (p *ParameterSet) indcpaKeyPair(rng io.Reader) (*indcpaPublicKey, *indcpaSecretKey, error) {
+	buf := make([]byte, SymSize+SymSize)
+	if _, err := io.ReadFull(rng, buf[:SymSize]); err != nil {
+		return nil, nil, err
+	}
+
+	sk := &indcpaSecretKey{
+		packed: make([]byte, p.indcpaSecretKeySize),
+	}
+	pk := &indcpaPublicKey{
+		packed: make([]byte, p.indcpaPublicKeySize),
+	}
+
+	h := sha3.New512()
+	h.Write(buf[:SymSize])
+	buf = buf[:0] // Reuse the backing store.
+	buf = h.Sum(buf)
+	publicSeed, noiseSeed := buf[:SymSize], buf[SymSize:]
+
+	a := p.allocMatrix()
+	genMatrix(a, publicSeed, false)
+
+	var nonce byte
+	skpv := p.allocPolyVec()
+	for _, pv := range skpv.vec {
+		pv.getNoise(noiseSeed, nonce, p.eta)
+		nonce++
+	}
+
+	skpv.ntt()
+
+	e := p.allocPolyVec()
+	for _, pv := range e.vec {
+		pv.getNoise(noiseSeed, nonce, p.eta)
+		nonce++
+	}
+
+	// matrix-vector multiplication
+	pkpv := p.allocPolyVec()
+	for i, pv := range pkpv.vec {
+		pv.pointwiseAcc(&skpv, &a[i])
+	}
+
+	pkpv.invntt()
+	pkpv.add(&pkpv, &e)
+
+	packSecretKey(sk.packed, &skpv)
+	packPublicKey(pk.packed, &pkpv, publicSeed)
+	pk.h = sha3.Sum256(pk.packed)
+
+	return pk, sk, nil
+}
+
+// Encryption function of the CPA-secure public-key encryption scheme
+// underlying Kyber.
+func (p *ParameterSet) indcpaEncrypt(c, m []byte, pk *indcpaPublicKey, coins []byte) {
+	var k, v, epp poly
+	var seed [SymSize]byte
+
+	pkpv := p.allocPolyVec()
+	unpackPublicKey(&pkpv, seed[:], pk.packed)
+
+	k.fromMsg(m)
+
+	pkpv.ntt()
+
+	at := p.allocMatrix()
+	genMatrix(at, seed[:], true)
+
+	var nonce byte
+	sp := p.allocPolyVec()
+	for _, pv := range sp.vec {
+		pv.getNoise(coins, nonce, p.eta)
+		nonce++
+	}
+
+	sp.ntt()
+
+	ep := p.allocPolyVec()
+	for _, pv := range ep.vec {
+		pv.getNoise(coins, nonce, p.eta)
+		nonce++
+	}
+
+	// matrix-vector multiplication
+	bp := p.allocPolyVec()
+	for i, pv := range bp.vec {
+		pv.pointwiseAcc(&sp, &at[i])
+	}
+
+	bp.invntt()
+	bp.add(&bp, &ep)
+
+	v.pointwiseAcc(&pkpv, &sp)
+	v.invntt()
+
+	epp.getNoise(coins, nonce, p.eta) // Don't need to increment nonce.
+
+	v.add(&v, &epp)
+	v.add(&v, &k)
+
+	packCiphertext(c, &bp, &v)
+}
+
+// Decryption function of the CPA-secure public-key encryption scheme
+// underlying Kyber.
+func (p *ParameterSet) indcpaDecrypt(m, c []byte, sk *indcpaSecretKey) {
+	var v, mp poly
+
+	skpv, bp := p.allocPolyVec(), p.allocPolyVec()
+	unpackCiphertext(&bp, &v, c)
+	unpackSecretKey(&skpv, sk.packed)
+
+	bp.ntt()
+
+	mp.pointwiseAcc(&skpv, &bp)
+	mp.invntt()
+
+	mp.sub(&mp, &v)
+
+	mp.toMsg(m)
+}
+
+func (p *ParameterSet) allocMatrix() []polyVec {
+	m := make([]polyVec, 0, p.k)
+	for i := 0; i < p.k; i++ {
+		m = append(m, p.allocPolyVec())
+	}
+	return m
+}
+
+func (p *ParameterSet) allocPolyVec() polyVec {
+	vec := make([]*poly, 0, p.k)
+	for i := 0; i < p.k; i++ {
+		vec = append(vec, new(poly))
+	}
+
+	return polyVec{vec}
+}

+ 183 - 0
kem.go

@@ -0,0 +1,183 @@
+// kem.go - Kyber key encapsulation mechanism.
+//
+// To the extent possible under law, Yawning Angel has waived all copyright
+// and related or neighboring rights to the software, using the Creative
+// Commons "CC0" public domain dedication. See LICENSE or
+// <http://creativecommons.org/publicdomain/zero/1.0/> for full details.
+
+package kyber
+
+import (
+	"bytes"
+	"crypto/subtle"
+	"errors"
+	"io"
+
+	"golang.org/x/crypto/sha3"
+)
+
+var (
+	// ErrInvalidKeySize is the error returned when a byte serailized key is
+	// an invalid size.
+	ErrInvalidKeySize = errors.New("kyber: invalid key size")
+
+	// ErrInvalidPrivateKey is the error returned when a byte serialized
+	// private key is malformed.
+	ErrInvalidPrivateKey = errors.New("kyber: invalid private key")
+)
+
+// PrivateKey is a Kyber private key.
+type PrivateKey struct {
+	PublicKey
+	sk *indcpaSecretKey
+	z  []byte
+}
+
+// Bytes returns the byte serialization of a PrivateKey.
+func (sk *PrivateKey) Bytes() []byte {
+	p := sk.PublicKey.p
+
+	b := make([]byte, 0, p.secretKeySize)
+	b = append(b, sk.sk.packed...)
+	b = append(b, sk.PublicKey.pk.packed...)
+	b = append(b, sk.PublicKey.pk.h[:]...)
+	b = append(b, sk.z...)
+
+	return b
+}
+
+// PrivateKeyFromBytes deserializes a byte serialized PrivateKey.
+func (p *ParameterSet) PrivateKeyFromBytes(b []byte) (*PrivateKey, error) {
+	if len(b) != p.secretKeySize {
+		return nil, ErrInvalidKeySize
+	}
+
+	sk := new(PrivateKey)
+	sk.sk = new(indcpaSecretKey)
+	sk.z = make([]byte, SymSize)
+	sk.PublicKey.pk = new(indcpaPublicKey)
+	sk.PublicKey.p = p
+
+	// De-serialize the public key first.
+	off := p.indcpaSecretKeySize
+	if err := sk.PublicKey.pk.fromBytes(p, b[off:off+p.publicKeySize]); err != nil {
+		return nil, err
+	}
+	off += p.publicKeySize
+	if !bytes.Equal(sk.PublicKey.pk.h[:], b[off:off+SymSize]) {
+		return nil, ErrInvalidPrivateKey
+	}
+	off += SymSize
+	copy(sk.z, b[off:])
+
+	// Then go back to de-serialize the private key.
+	if err := sk.sk.fromBytes(p, b[:p.indcpaSecretKeySize]); err != nil {
+		return nil, err
+	}
+
+	return sk, nil
+}
+
+// PublicKey is a Kyber public key.
+type PublicKey struct {
+	pk *indcpaPublicKey
+	p  *ParameterSet
+}
+
+// Bytes returns the byte serialization of a PublicKey.
+func (pk *PublicKey) Bytes() []byte {
+	return pk.pk.toBytes()
+}
+
+// PublicKeyFromBytes deserializes a byte serialized PublicKey.
+func (p *ParameterSet) PublicKeyFromBytes(b []byte) (*PublicKey, error) {
+	pk := &PublicKey{
+		pk: new(indcpaPublicKey),
+		p:  p,
+	}
+
+	if err := pk.pk.fromBytes(p, b); err != nil {
+		return nil, err
+	}
+
+	return pk, nil
+}
+
+// GenerateKeyPair generates a private and public key parameterized with the
+// given ParameterSet.
+func (p *ParameterSet) GenerateKeyPair(rng io.Reader) (*PublicKey, *PrivateKey, error) {
+	kp := new(PrivateKey)
+
+	var err error
+	if kp.PublicKey.pk, kp.sk, err = p.indcpaKeyPair(rng); err != nil {
+		return nil, nil, err
+	}
+
+	kp.PublicKey.p = p
+	kp.z = make([]byte, SymSize)
+	if _, err := io.ReadFull(rng, kp.z); err != nil {
+		return nil, nil, err
+	}
+
+	return &kp.PublicKey, kp, nil
+}
+
+// KEMEncrypt generates cipher text and shared secret via the CCA-secure Kyber
+// key encapsulation mechanism.
+func (pk *PublicKey) KEMEncrypt(rng io.Reader) (cipherText []byte, sharedSecret []byte, err error) {
+	var buf [SymSize]byte
+	if _, err = io.ReadFull(rng, buf[:]); err != nil {
+		return
+	}
+	buf = sha3.Sum256(buf[:]) // Don't release system RNG output
+
+	hKr := sha3.New512()
+	hKr.Write(buf[:])
+	hKr.Write(pk.pk.h[:]) // Multitarget countermeasures for coins + contributory KEM
+	kr := hKr.Sum(nil)
+
+	cipherText = make([]byte, pk.p.cipherTextSize)
+	pk.p.indcpaEncrypt(cipherText, buf[:], pk.pk, kr[SymSize:]) // coins are in kr[SymSize:]
+
+	hc := sha3.Sum256(cipherText)
+	copy(kr[SymSize:], hc[:]) // overwrite coins in kr with H(c)
+	hSs := sha3.New256()
+	hSs.Write(kr)
+	sharedSecret = hSs.Sum(nil) // hash concatenation of pre-k and H(c) to k
+
+	return
+}
+
+// KEMDecrypt generates shared secret for given cipher text via the CCA-secure
+// Kyber key encapsulation mechanism.
+//
+// On success fail will be 0, otherwise fail will be set to -1 and
+// sharedSecret will contain a randomized value.
+func (sk *PrivateKey) KEMDecrypt(cipherText []byte) (sharedSecret []byte, fail int) {
+	var buf [2 * SymSize]byte
+
+	p := sk.PublicKey.p
+	if len(cipherText) != p.CipherTextSize() {
+		return nil, -1
+	}
+	p.indcpaDecrypt(buf[:SymSize], cipherText, sk.sk)
+
+	copy(buf[SymSize:], sk.PublicKey.pk.h[:]) // Multitarget countermeasure for coins + contributory KEM
+	kr := sha3.Sum512(buf[:])
+
+	cmp := make([]byte, p.cipherTextSize)
+	p.indcpaEncrypt(cmp, buf[:SymSize], sk.PublicKey.pk, kr[SymSize:]) // coins are in kr[SymSize:]
+
+	hc := sha3.Sum256(cipherText)
+	copy(kr[SymSize:], hc[:]) // overwrite coins in kr with H(c)
+
+	fail = subtle.ConstantTimeSelect(subtle.ConstantTimeCompare(cipherText, cmp), 0, 1)
+	subtle.ConstantTimeCopy(fail, kr[SymSize:], sk.z) // Overwrite pre-k with z on re-encryption failure
+	fail = -fail
+
+	h := sha3.New256()
+	h.Write(kr[:])
+	sharedSecret = h.Sum(nil)
+
+	return
+}

+ 185 - 0
kem_test.go

@@ -0,0 +1,185 @@
+// kem_test.go - Kyber KEM tests.
+//
+// To the extent possible under law, Yawning Angel has waived all copyright
+// and related or neighboring rights to the software, using the Creative
+// Commons "CC0" public domain dedication. See LICENSE or
+// <http://creativecommons.org/publicdomain/zero/1.0/> for full details.
+
+package kyber
+
+import (
+	"bytes"
+	"crypto/rand"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+const nTests = 100
+
+var allParams = []*ParameterSet{
+	Kyber512,
+	Kyber768,
+	Kyber1024,
+}
+
+func TestKEM(t *testing.T) {
+	for _, p := range allParams {
+		t.Run(p.Name()+"_Keys", func(t *testing.T) { doTestKEMKeys(t, p) })
+		t.Run(p.Name()+"_Invalid_SecretKey_A", func(t *testing.T) { doTestKEMInvalidSkA(t, p) })
+		t.Run(p.Name()+"_Invalid_CipherText", func(t *testing.T) { doTestKEMInvalidCipherText(t, p) })
+	}
+}
+
+func doTestKEMKeys(t *testing.T, p *ParameterSet) {
+	require := require.New(t)
+
+	t.Logf("PrivateKeySize(): %v", p.PrivateKeySize())
+	t.Logf("PublicKeySize(): %v", p.PublicKeySize())
+	t.Logf("CipherTextSize(): %v", p.CipherTextSize())
+
+	for i := 0; i < nTests; i++ {
+		// Generate a key pair.
+		pk, sk, err := p.GenerateKeyPair(rand.Reader)
+		require.NoError(err, "GenerateKeyPair()")
+
+		// Test serialization.
+		b := sk.Bytes()
+		require.Len(b, p.PrivateKeySize(), "sk.Bytes(): Length")
+		sk2, err := p.PrivateKeyFromBytes(b)
+		require.NoError(err, "PrivateKeyFromBytes(b)")
+		requirePrivateKeyEqual(require, sk, sk2)
+
+		b = pk.Bytes()
+		require.Len(b, p.PublicKeySize(), "pk.Bytes(): Length")
+		pk2, err := p.PublicKeyFromBytes(b)
+		require.NoError(err, "PublicKeyFromBytes(b)")
+		requirePublicKeyEqual(require, pk, pk2)
+
+		// Test encrypt/decrypt.
+		ct, ss, err := pk.KEMEncrypt(rand.Reader)
+		require.NoError(err, "KEMEncrypt()")
+		require.Len(ct, p.CipherTextSize(), "KEMEncrypt(): ct Length")
+		require.Len(ss, SymSize, "KEMEncrypt(): ss Length")
+
+		ss2, fail := sk.KEMDecrypt(ct)
+		require.Equal(0, fail, "KEMDecrypt(): fail")
+		require.Equal(ss, ss2, "KEMDecrypt(): ss")
+	}
+}
+
+func doTestKEMInvalidSkA(t *testing.T, p *ParameterSet) {
+	require := require.New(t)
+
+	for i := 0; i < nTests; i++ {
+		// Alice generates a public key.
+		pk, skA, err := p.GenerateKeyPair(rand.Reader)
+		require.NoError(err, "GenerateKeyPair()")
+
+		// Bob derives a secret key and creates a response.
+		sendB, keyB, err := pk.KEMEncrypt(rand.Reader)
+		require.NoError(err, "KEMEncrypt()")
+
+		// Replace secret key with random values.
+		_, err = rand.Read(skA.sk.packed)
+		require.NoError(err, "rand.Read()")
+
+		// Alice uses Bob's response to get her secret key.
+		keyA, fail := skA.KEMDecrypt(sendB)
+		require.Equal(-1, fail, "KEMDecrypt(): fail")
+		require.NotEqual(keyA, keyB, "KEMDecrypt(): ss")
+	}
+}
+
+func doTestKEMInvalidCipherText(t *testing.T, p *ParameterSet) {
+	require := require.New(t)
+	var rawPos [2]byte
+
+	ciphertextSize := p.CipherTextSize()
+
+	for i := 0; i < nTests; i++ {
+		_, err := rand.Read(rawPos[:])
+		require.NoError(err, "rand.Read()")
+		pos := (int(rawPos[0]) << 8) | int(rawPos[1])
+
+		// Alice generates a public key.
+		pk, skA, err := p.GenerateKeyPair(rand.Reader)
+		require.NoError(err, "GenerateKeyPair()")
+
+		// Bob derives a secret key and creates a response.
+		sendB, keyB, err := pk.KEMEncrypt(rand.Reader)
+		require.NoError(err, "KEMEncrypt()")
+
+		// Change some byte in the ciphertext (i.e., encapsulated key).
+		sendB[pos%ciphertextSize] ^= 23
+
+		// Alice uses Bob's response to get her secret key.
+		keyA, fail := skA.KEMDecrypt(sendB)
+		require.Equal(-1, fail, "KEMDecrypt(): fail")
+		require.NotEqual(keyA, keyB, "KEMDecrypt(): ss")
+	}
+}
+
+func requirePrivateKeyEqual(require *require.Assertions, a, b *PrivateKey) {
+	require.EqualValues(a.sk, b.sk, "sk (indcpaSecretKey)")
+	require.Equal(a.z, b.z, "z (random bytes)")
+	requirePublicKeyEqual(require, &a.PublicKey, &b.PublicKey)
+}
+
+func requirePublicKeyEqual(require *require.Assertions, a, b *PublicKey) {
+	require.EqualValues(a.pk, b.pk, "pk (indcpaPublicKey)")
+	require.Equal(a.p, b.p, "p (ParameterSet)")
+}
+
+func BenchmarkKEM(b *testing.B) {
+	for _, p := range allParams {
+		b.Run(p.Name()+"_GenerateKeyPair", func(b *testing.B) { doBenchKEMGenerateKeyPair(b, p) })
+		b.Run(p.Name()+"_KEMEncrypt", func(b *testing.B) { doBenchKEMEncDec(b, p, true) })
+		b.Run(p.Name()+"_KEMDecrypt", func(b *testing.B) { doBenchKEMEncDec(b, p, false) })
+	}
+}
+
+func doBenchKEMGenerateKeyPair(b *testing.B, p *ParameterSet) {
+	for i := 0; i < b.N; i++ {
+		_, _, err := p.GenerateKeyPair(rand.Reader)
+		if err != nil {
+			b.Fatalf("GenerateKeyPair(): %v", err)
+		}
+	}
+}
+
+func doBenchKEMEncDec(b *testing.B, p *ParameterSet, isEnc bool) {
+	b.StopTimer()
+	for i := 0; i < b.N; i++ {
+		pk, skA, err := p.GenerateKeyPair(rand.Reader)
+		if err != nil {
+			b.Fatalf("GenerateKeyPair(): %v", err)
+		}
+
+		if isEnc {
+			b.StartTimer()
+		}
+
+		sendB, keyB, err := pk.KEMEncrypt(rand.Reader)
+		if err != nil {
+			b.Fatalf("KEMEncrypt(): %v", err)
+		}
+		if isEnc {
+			b.StopTimer()
+		} else {
+			b.StartTimer()
+		}
+
+		keyA, fail := skA.KEMDecrypt(sendB)
+		if !isEnc {
+			b.StopTimer()
+		}
+
+		if fail != 0 {
+			b.Fatalf("KEMDecrypt(): fail %v", fail)
+		}
+		if !bytes.Equal(keyA, keyB) {
+			b.Fatalf("KEMDecrypt(): key mismatch")
+		}
+	}
+}

+ 294 - 0
kem_vectors_test.go

@@ -0,0 +1,294 @@
+// kem_vectors_test.go - Kyber KEM test vector tests.
+//
+// To the extent possible under law, Yawning Angel has waived all copyright
+// and related or neighboring rights to the software, using the Creative
+// Commons "CC0" public domain dedication. See LICENSE or
+// <http://creativecommons.org/publicdomain/zero/1.0/> for full details.
+
+package kyber
+
+import (
+	"bufio"
+	"crypto/sha256"
+	"encoding/hex"
+	"encoding/json"
+	"errors"
+	"io"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+const nrTestVectors = 1000 // WARNING: Must match the reference code.
+
+var compactTestVectors = make(map[string][]byte)
+
+func TestKEMVectors(t *testing.T) {
+	if err := loadCompactTestVectors(); err != nil {
+		t.Fatalf("loadCompactTestVectors(): %v", err)
+	}
+
+	for _, p := range allParams {
+		t.Run(p.Name(), func(t *testing.T) { doTestKEMVectors(t, p) })
+	}
+}
+
+func doTestKEMVectors(t *testing.T, p *ParameterSet) {
+	require := require.New(t)
+
+	// The full test vectors are gigantic, and aren't checked into the
+	// git repository.
+	vecs, err := loadTestVectors(p)
+	if err == nil {
+		// If they exist because someone generated them and placed them in
+		// the correct location, use them.
+		doTestKEMVectorsFull(require, p, vecs)
+	} else {
+		// Otherwise use the space saving representation based on comparing
+		// digests.
+		doTestKEMVectorsCompact(require, p)
+	}
+}
+
+func doTestKEMVectorsFull(require *require.Assertions, p *ParameterSet, vecs []*vector) {
+	rng := newTestRng()
+	for idx, vec := range vecs {
+		pk, sk, err := p.GenerateKeyPair(rng)
+		require.NoError(err, "GenerateKeyPair(): %v", idx)
+		require.Equal(vec.rndKP, rng.PopHist(), "randombytes() kp: %v", idx)
+		require.Equal(vec.rndZ, rng.PopHist(), "randombytes() z: %v", idx)
+		require.Equal(vec.pk, pk.Bytes(), "pk: %v", idx)
+		require.Equal(vec.skA, sk.Bytes(), "skA: %v", idx)
+
+		sendB, keyB, err := pk.KEMEncrypt(rng)
+		require.NoError(err, "KEMEncrypt(): %v", idx)
+		require.Equal(vec.rndEnc, rng.PopHist(), "randombytes() enc: %v", idx)
+		require.Equal(vec.sendB, sendB, "sendB: %v", idx)
+		require.Equal(vec.keyB, keyB, "keyB: %v", idx)
+
+		keyA, fail := sk.KEMDecrypt(sendB)
+		require.Equal(0, fail, "fail: %v", idx)
+		require.Equal(vec.keyA, keyA, "keyA: %v", idx)
+	}
+}
+
+func doTestKEMVectorsCompact(require *require.Assertions, p *ParameterSet) {
+	h := sha256.New()
+
+	rng := newTestRng()
+	for idx := 0; idx < nrTestVectors; idx++ {
+		pk, sk, err := p.GenerateKeyPair(rng)
+		require.NoError(err, "GenerateKeyPair(): %v", idx)
+		h.Write([]byte(hex.EncodeToString(rng.PopHist()) + "\n"))
+		h.Write([]byte(hex.EncodeToString(rng.PopHist()) + "\n"))
+		h.Write([]byte(hex.EncodeToString(pk.Bytes()) + "\n"))
+		h.Write([]byte(hex.EncodeToString(sk.Bytes()) + "\n"))
+
+		sendB, keyB, err := pk.KEMEncrypt(rng)
+		require.NoError(err, "KEMEncrypt(): %v", idx)
+		h.Write([]byte(hex.EncodeToString(rng.PopHist()) + "\n"))
+		h.Write([]byte(hex.EncodeToString(sendB) + "\n"))
+		h.Write([]byte(hex.EncodeToString(keyB) + "\n"))
+
+		keyA, fail := sk.KEMDecrypt(sendB)
+		require.Equal(0, fail, "fail: %v", idx)
+		h.Write([]byte(hex.EncodeToString(keyA) + "\n"))
+	}
+
+	require.Equal(compactTestVectors[p.Name()], h.Sum(nil), "Digest mismatch")
+}
+
+func loadCompactTestVectors() error {
+	f, err := os.Open(filepath.Join("testdata", "compactVectors.json"))
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	rawMap := make(map[string]string)
+	dec := json.NewDecoder(f)
+	if err = dec.Decode(&rawMap); err != nil {
+		return err
+	}
+
+	for k, v := range rawMap {
+		digest, err := hex.DecodeString(v)
+		if err != nil {
+			return err
+		}
+
+		compactTestVectors[k] = digest
+	}
+
+	return nil
+}
+
+type vector struct {
+	rndKP  []byte
+	rndZ   []byte
+	pk     []byte
+	skA    []byte
+	rndEnc []byte
+	sendB  []byte
+	keyB   []byte
+	keyA   []byte
+}
+
+func loadTestVectors(p *ParameterSet) ([]*vector, error) {
+	fn := "KEM-" + p.Name() + ".full"
+
+	f, err := os.Open(filepath.Join("testdata", fn))
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+
+	var vectors []*vector
+	scanner := bufio.NewScanner(f)
+	for {
+		v, err := getNextVector(scanner)
+		switch err {
+		case nil:
+			vectors = append(vectors, v)
+		case io.EOF:
+			return vectors, nil
+		default:
+			return nil, err
+		}
+	}
+}
+
+func getNextVector(scanner *bufio.Scanner) (*vector, error) {
+	var v [][]byte
+
+	for i := 0; i < 8; i++ {
+		if ok := scanner.Scan(); !ok {
+			if i == 0 {
+				return nil, io.EOF
+			}
+			return nil, errors.New("truncated file")
+		}
+		b, err := hex.DecodeString(scanner.Text())
+		if err != nil {
+			return nil, err
+		}
+		v = append(v, b)
+	}
+
+	vec := &vector{
+		rndKP:  v[0],
+		rndZ:   v[1],
+		pk:     v[2],
+		skA:    v[3],
+		rndEnc: v[4],
+		sendB:  v[5],
+		keyB:   v[6],
+		keyA:   v[7],
+	}
+
+	return vec, nil
+}
+
+type testRNG struct {
+	seed    [32]uint32
+	in      [12]uint32
+	out     [8]uint32
+	outleft int
+
+	hist [][]byte
+}
+
+func newTestRng() *testRNG {
+	r := new(testRNG)
+	r.seed = [32]uint32{
+		3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 9, 3, 2, 3, 8, 4, 6, 2, 6, 4, 3, 3, 8, 3, 2, 7, 9, 5,
+	}
+	for i := range r.in {
+		r.in[i] = 0
+	}
+	r.outleft = 0
+	return r
+}
+
+func (r *testRNG) surf() {
+	var t [12]uint32
+	var sum uint32
+
+	for i, v := range r.in {
+		t[i] = v ^ r.seed[12+i]
+	}
+	for i := range r.out {
+		r.out[i] = r.seed[24+i]
+	}
+	x := t[11]
+	rotate := func(x uint32, b uint) uint32 {
+		return (((x) << (b)) | ((x) >> (32 - (b))))
+	}
+	mush := func(i int, b uint) {
+		t[i] += (((x ^ r.seed[i]) + sum) ^ rotate(x, b))
+		x = t[i]
+	}
+	for loop := 0; loop < 2; loop++ {
+		for rr := 0; rr < 16; rr++ {
+			sum += 0x9e3779b9
+			mush(0, 5)
+			mush(1, 7)
+			mush(2, 9)
+			mush(3, 13)
+			mush(4, 5)
+			mush(5, 7)
+			mush(6, 9)
+			mush(7, 13)
+			mush(8, 5)
+			mush(9, 7)
+			mush(10, 9)
+			mush(11, 13)
+		}
+		for i := range r.out {
+			r.out[i] ^= t[i+4]
+		}
+	}
+}
+
+func (r *testRNG) Read(x []byte) (n int, err error) {
+	dst := x
+
+	xlen, ret := len(x), len(x)
+	for xlen > 0 {
+		if r.outleft == 0 {
+			r.in[0]++
+			if r.in[0] == 0 {
+				r.in[1]++
+				if r.in[1] == 0 {
+					r.in[2]++
+					if r.in[2] == 0 {
+						r.in[3]++
+					}
+				}
+			}
+			r.surf()
+			r.outleft = 8
+		}
+		r.outleft--
+		x[0] = byte(r.out[r.outleft])
+		x = x[1:]
+		xlen--
+	}
+
+	r.hist = append(r.hist, append([]byte{}, dst...))
+
+	return ret, nil
+}
+
+func (r *testRNG) PopHist() []byte {
+	if len(r.hist) == 0 {
+		panic("pop underflow")
+	}
+
+	b := r.hist[0]
+	r.hist = append([][]byte{}, r.hist[1:]...)
+
+	return b
+}

+ 241 - 0
kex.go

@@ -0,0 +1,241 @@
+// kex.go - Kyber key exchange.
+//
+// To the extent possible under law, Yawning Angel has waived all copyright
+// and related or neighboring rights to the software, using the Creative
+// Commons "CC0" public domain dedication. See LICENSE or
+// <http://creativecommons.org/publicdomain/zero/1.0/> for full details.
+
+package kyber
+
+import (
+	"io"
+
+	"golang.org/x/crypto/sha3"
+)
+
+// UAKEInitiatorMessageSize returns the size of the initiator UAKE message
+// in bytes.
+func (p *ParameterSet) UAKEInitiatorMessageSize() int {
+	return p.PublicKeySize() + p.CipherTextSize()
+}
+
+// UAKEResponderMessageSize returns the size of the responder UAKE message
+// in bytes.
+func (p *ParameterSet) UAKEResponderMessageSize() int {
+	return p.CipherTextSize()
+}
+
+// UAKEInitiatorState is a initiator UAKE instance.  Each instance MUST only
+// be used for one key exchange and never reused.
+type UAKEInitiatorState struct {
+	// Message is the UAKE message to send to the responder.
+	Message []byte
+
+	eSk *PrivateKey
+	tk  []byte
+}
+
+// Shared generates a shared secret for the given UAKE instance and responder
+// message.
+//
+// On success fail will be 0, otherwise fail will be set to -1 and
+// sharedSecret will contain a randomized value.
+func (s *UAKEInitiatorState) Shared(recv []byte) (sharedSecret []byte, fail int) {
+	xof := sha3.NewShake256()
+	var tk []byte
+
+	tk, fail = s.eSk.KEMDecrypt(recv)
+	xof.Write(tk)
+	xof.Write(s.tk)
+	sharedSecret = make([]byte, SymSize)
+	xof.Read(sharedSecret)
+
+	return
+}
+
+// NewUAKEInitiatorState creates a new initiator UAKE instance.
+func (pk *PublicKey) NewUAKEInitiatorState(rng io.Reader) (*UAKEInitiatorState, error) {
+	s := new(UAKEInitiatorState)
+	s.Message = make([]byte, 0, pk.p.UAKEInitiatorMessageSize())
+
+	var err error
+	_, s.eSk, err = pk.p.GenerateKeyPair(rng)
+	if err != nil {
+		return nil, err
+	}
+	s.Message = append(s.Message, s.eSk.PublicKey.Bytes()...)
+
+	var ct []byte
+	ct, s.tk, err = pk.KEMEncrypt(rng)
+	if err != nil {
+		return nil, err
+	}
+
+	s.Message = append(s.Message, ct...)
+
+	return s, nil
+}
+
+// UAKEResponderShared generates a responder message and shared secret given
+// a initiator UAKE message.
+//
+// On success fail will be 0, otherwise fail will be set to -1 and
+// sharedSecret will contain a randomized value.
+func (sk *PrivateKey) UAKEResponderShared(rng io.Reader, recv []byte) (message, sharedSecret []byte, fail int) {
+	p := sk.PublicKey.p
+	pkLen := p.PublicKeySize()
+
+	fail = -1
+
+	// Deserialize the peer's ephemeral public key.
+	if len(recv) != p.UAKEInitiatorMessageSize() {
+		return
+	}
+	rawPk, ct := recv[:pkLen], recv[pkLen:]
+	pk, err := p.PublicKeyFromBytes(rawPk)
+	if err != nil {
+		return
+	}
+
+	xof := sha3.NewShake256()
+	var tk []byte
+
+	message, tk, err = pk.KEMEncrypt(rng)
+	if err != nil {
+		return
+	}
+	xof.Write(tk)
+
+	tk, fail = sk.KEMDecrypt(ct)
+	xof.Write(tk)
+	sharedSecret = make([]byte, SymSize)
+	xof.Read(sharedSecret)
+
+	return
+}
+
+// AKEInitiatorMessageSize returns the size of the initiator AKE message
+// in bytes.
+func (p *ParameterSet) AKEInitiatorMessageSize() int {
+	return p.PublicKeySize() + p.CipherTextSize()
+}
+
+// AKEResponderMessageSize returns the size of the responder AKE message
+// in bytes.
+func (p *ParameterSet) AKEResponderMessageSize() int {
+	return 2 * p.CipherTextSize()
+}
+
+// AKEInitiatorState is a initiator AKE instance.  Each instance MUST only be
+// used for one key exchange and never reused.
+type AKEInitiatorState struct {
+	// Message is the AKE message to send to the responder.
+	Message []byte
+
+	eSk *PrivateKey
+	tk  []byte
+}
+
+// Shared generates a shared secret for the given AKE instance, responder
+// message, and long term initiator private key.
+//
+// On success fail will be 0, otherwise fail will be set to -1 and
+// sharedSecret will contain a randomized value.
+func (s *AKEInitiatorState) Shared(recv []byte, initiatorPrivateKey *PrivateKey) (sharedSecret []byte, fail int) {
+	p := s.eSk.PublicKey.p
+	fail = -1
+
+	if initiatorPrivateKey.PublicKey.p != p {
+		return
+	}
+	if len(recv) != p.AKEResponderMessageSize() {
+		return
+	}
+	ctLen := p.CipherTextSize()
+
+	xof := sha3.NewShake256()
+	var tk []byte
+
+	tk, fail = s.eSk.KEMDecrypt(recv[:ctLen])
+	xof.Write(tk)
+
+	var fail2 int
+	tk, fail2 = initiatorPrivateKey.KEMDecrypt(recv[ctLen:])
+	xof.Write(tk)
+	fail |= fail2
+
+	xof.Write(s.tk)
+	sharedSecret = make([]byte, SymSize)
+	xof.Read(sharedSecret)
+
+	return
+}
+
+// NewAKEInitiatorState creates a new initiator AKE instance.
+func (pk *PublicKey) NewAKEInitiatorState(rng io.Reader) (*AKEInitiatorState, error) {
+	s := new(AKEInitiatorState)
+
+	// This is identical to the UAKE case, so just reuse the code.
+	us, err := pk.NewUAKEInitiatorState(rng)
+	if err != nil {
+		return nil, err
+	}
+
+	s.Message = us.Message
+	s.eSk = us.eSk
+	s.tk = us.tk
+
+	return s, nil
+}
+
+// AKEResponderShared generates a responder message and shared secret given
+// a initiator AKE message and long term initiator public key.
+//
+// On success fail will be 0, otherwise fail will be set to -1 and
+// sharedSecret will contain a randomized value.
+func (sk *PrivateKey) AKEResponderShared(rng io.Reader, recv []byte, peerPublicKey *PublicKey) (message, sharedSecret []byte, fail int) {
+	p := sk.PublicKey.p
+	pkLen := p.PublicKeySize()
+
+	fail = -1
+
+	if peerPublicKey.p != p {
+		return
+	}
+
+	// Deserialize the peer's ephemeral public key.
+	if len(recv) != p.AKEInitiatorMessageSize() {
+		return
+	}
+	rawPk, ct := recv[:pkLen], recv[pkLen:]
+	pk, err := p.PublicKeyFromBytes(rawPk)
+	if err != nil {
+		return
+	}
+
+	message = make([]byte, 0, p.AKEResponderMessageSize())
+
+	xof := sha3.NewShake256()
+	var tk, tmp []byte
+
+	tmp, tk, err = pk.KEMEncrypt(rng)
+	if err != nil {
+		return
+	}
+	xof.Write(tk)
+	message = append(message, tmp...)
+
+	tmp, tk, err = peerPublicKey.KEMEncrypt(rng)
+	if err != nil {
+		return
+	}
+	xof.Write(tk)
+	message = append(message, tmp...)
+
+	tk, fail = sk.KEMDecrypt(ct)
+	xof.Write(tk)
+	sharedSecret = make([]byte, SymSize)
+	xof.Read(sharedSecret)
+
+	return
+}

+ 83 - 0
kex_test.go

@@ -0,0 +1,83 @@
+// kex_test.go - Kyber key exchange tests.
+//
+// To the extent possible under law, Yawning Angel has waived all copyright
+// and related or neighboring rights to the software, using the Creative
+// Commons "CC0" public domain dedication. See LICENSE or
+// <http://creativecommons.org/publicdomain/zero/1.0/> for full details.
+
+package kyber
+
+import (
+	"crypto/rand"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestAKE(t *testing.T) {
+	for _, p := range allParams {
+		t.Run(p.Name()+"_UAKE", func(t *testing.T) { doTestUAKE(t, p) })
+		t.Run(p.Name()+"_AKE", func(t *testing.T) { doTestAKE(t, p) })
+	}
+}
+
+func doTestUAKE(t *testing.T, p *ParameterSet) {
+	require := require.New(t)
+
+	t.Logf("UAKEInitiatorMessageSize(): %v", p.UAKEInitiatorMessageSize())
+	t.Logf("UAKEResponderMessageSize(): %v", p.UAKEResponderMessageSize())
+
+	for i := 0; i < nTests; i++ {
+		// Generate the responder key pair.
+		pkB, skB, err := p.GenerateKeyPair(rand.Reader)
+		require.NoError(err, "GenerateKeyPair()")
+
+		// Create the initiator state.
+		stateA, err := pkB.NewUAKEInitiatorState(rand.Reader)
+		require.NoError(err, "NewUAKEInitiatorState()")
+		require.Len(stateA.Message, p.UAKEInitiatorMessageSize(), "stateA.Message: Length")
+
+		// Create the responder message and shared secret.
+		msgB, ssB, fail := skB.UAKEResponderShared(rand.Reader, stateA.Message)
+		require.Equal(0, fail, "UAKEResponderShared(): fail")
+		require.Len(msgB, p.UAKEResponderMessageSize(), "UAKEResponderShared(): msgB Length")
+		require.Len(ssB, SymSize, "UAKEResponderShared(): ssB Length")
+
+		// Create the initiator shared secret.
+		ssA, fail := stateA.Shared(msgB)
+		require.Equal(0, fail, "stateA.Shared(): fail")
+		require.Equal(ssA, ssB, "Shared secret mismatch")
+	}
+}
+
+func doTestAKE(t *testing.T, p *ParameterSet) {
+	require := require.New(t)
+
+	t.Logf("AKEInitiatorMessageSize(): %v", p.AKEInitiatorMessageSize())
+	t.Logf("AKEResponderMessageSize(): %v", p.AKEResponderMessageSize())
+
+	for i := 0; i < nTests; i++ {
+		// Generate the initiator and responder key pairs.
+		pkB, skB, err := p.GenerateKeyPair(rand.Reader)
+		require.NoError(err, "GenerateKeyPair(): Responder")
+
+		pkA, skA, err := p.GenerateKeyPair(rand.Reader)
+		require.NoError(err, "GenerateKeyPair(): Initiator")
+
+		// Create the initiator state.
+		stateA, err := pkB.NewAKEInitiatorState(rand.Reader)
+		require.NoError(err, "NewAKEInitiatorState()")
+		require.Len(stateA.Message, p.AKEInitiatorMessageSize(), "stateA.Message: Length")
+
+		// Create the responder message and shared secret.
+		msgB, ssB, fail := skB.AKEResponderShared(rand.Reader, stateA.Message, pkA)
+		require.Equal(0, fail, "AKEResponderShared(): fail")
+		require.Len(msgB, p.AKEResponderMessageSize(), "AKEResponderShared(): msgB Length")
+		require.Len(ssB, SymSize, "AKEResponderShared(): ssB Length")
+
+		// Create the initiator shared secret.
+		ssA, fail := stateA.Shared(msgB, skA)
+		require.Equal(0, fail, "stateA.Shared(): fail")
+		require.Equal(ssA, ssB, "Shared secret mismatch")
+	}
+}

+ 65 - 0
ntt.go

@@ -0,0 +1,65 @@
+// ntt.go - Number-Theoretic Transform.
+//
+// To the extent possible under law, Yawning Angel has waived all copyright
+// and related or neighboring rights to the software, using the Creative
+// Commons "CC0" public domain dedication. See LICENSE or
+// <http://creativecommons.org/publicdomain/zero/1.0/> for full details.
+
+package kyber
+
+// Computes negacyclic number-theoretic transform (NTT) of a polynomial (vector
+// of 256 coefficients) in place; inputs assumed to be in normal order, output
+// in bitreversed order.
+func ntt(p *[kyberN]uint16) {
+	var j int
+	k := 1
+	for level := 7; level >= 0; level-- {
+		distance := 1 << uint(level)
+		for start := 0; start < kyberN; start = j + distance {
+			zeta := zetas[k]
+			k++
+			for j = start; j < start+distance; j++ {
+				t := montgomeryReduce(uint32(zeta) * uint32(p[j+distance]))
+				p[j+distance] = barrettReduce(p[j] + 4*kyberQ - t)
+
+				if level&1 == 1 { // odd level
+					p[j] = p[j] + t // Omit reduction (be lazy)
+				} else {
+					p[j] = barrettReduce(p[j] + t)
+				}
+			}
+		}
+	}
+}
+
+// Computes inverse of negacyclic number-theoretic transform (NTT) of a
+// polynomial (vector of 256 coefficients) in place; inputs assumed to be in
+// bitreversed order, output in normal order.
+func invntt(a *[kyberN]uint16) {
+	for level := 0; level < 8; level++ {
+		distance := 1 << uint(level)
+		for start := 0; start < distance; start++ {
+			var jTwiddle int
+			for j := start; j < kyberN-1; j += 2 * distance {
+				w := uint32(omegasInvBitrevMontgomery[jTwiddle])
+				jTwiddle++
+
+				temp := a[j]
+
+				if level&1 == 1 { // odd level
+					a[j] = barrettReduce(temp + a[j+distance])
+				} else {
+					a[j] = temp + a[j+distance] // Omit reduction (be lazy)
+				}
+
+				t := w * (uint32(temp) + 4*kyberQ - uint32(a[j+distance]))
+
+				a[j+distance] = montgomeryReduce(t)
+			}
+		}
+	}
+
+	for i, v := range psisInvMontgomery {
+		a[i] = montgomeryReduce(uint32(a[i]) * uint32(v))
+	}
+}

+ 116 - 0
params.go

@@ -0,0 +1,116 @@
+// params.go - Kyber parameterization.
+//
+// To the extent possible under law, Yawning Angel has waived all copyright
+// and related or neighboring rights to the software, using the Creative
+// Commons "CC0" public domain dedication. See LICENSE or
+// <http://creativecommons.org/publicdomain/zero/1.0/> for full details.
+
+package kyber
+
+const (
+	// SymSize is the size of the shared key (and certain internal parameters
+	// such as hashes and seeds) in bytes.
+	SymSize = 32
+
+	kyberN = 256
+	kyberQ = 7681
+
+	polySize           = 416
+	polyCompressedSize = 96
+
+	compressedCoeffSize = 352
+)
+
+var (
+	// Kyber512 is the Kyber-512 parameter set, which aims to provide security
+	// equivalent to AES-128.
+	//
+	// This parameter set has a 1632 byte private key, 736 byte public key,
+	// and a 800 byte cipher text.
+	Kyber512 = newParameterSet("Kyber-512", 2)
+
+	// Kyber768 is the Kyber-768 parameter set, which aims to provide security
+	// equivalent to AES-192.
+	//
+	// This parameter set has a 2400 byte private key, 1088 byte public key,
+	// and a 1088 byte cipher text.
+	Kyber768 = newParameterSet("Kyber-768", 3)
+
+	// Kyber1024 is the Kyber-1024 parameter set, which aims to provide
+	// security equivalent to AES-256.
+	//
+	// This parameter set has a 3168 byte private key, 1440 byte public key,
+	// and a 1504 byte cipher text.
+	Kyber1024 = newParameterSet("Kyber-1024", 4)
+)
+
+// ParameterSet is a Kyber parameter set.
+type ParameterSet struct {
+	name string
+
+	k   int
+	eta int
+
+	polyVecSize           int
+	polyVecCompressedSize int
+
+	indcpaMsgSize       int
+	indcpaPublicKeySize int
+	indcpaSecretKeySize int
+	indcpaSize          int
+
+	publicKeySize  int
+	secretKeySize  int
+	cipherTextSize int
+}
+
+// Name returns the name of a given ParameterSet.
+func (p *ParameterSet) Name() string {
+	return p.name
+}
+
+// PublicKeySize returns the size of a public key in bytes.
+func (p *ParameterSet) PublicKeySize() int {
+	return p.publicKeySize
+}
+
+// PrivateKeySize returns the size of a private key in bytes.
+func (p *ParameterSet) PrivateKeySize() int {
+	return p.secretKeySize
+}
+
+// CipherTextSize returns the size of a cipher text in bytes.
+func (p *ParameterSet) CipherTextSize() int {
+	return p.cipherTextSize
+}
+
+func newParameterSet(name string, k int) *ParameterSet {
+	var p ParameterSet
+
+	p.name = name
+	p.k = k
+	switch k {
+	case 2:
+		p.eta = 5
+	case 3:
+		p.eta = 4
+	case 4:
+		p.eta = 3
+	default:
+		panic("kyber: k must be in {2,3,4}")
+	}
+
+	p.polyVecSize = k * polySize
+	p.polyVecCompressedSize = k * compressedCoeffSize
+
+	p.indcpaMsgSize = SymSize
+	p.indcpaPublicKeySize = p.polyVecCompressedSize + SymSize
+	p.indcpaSecretKeySize = p.polyVecSize
+	p.indcpaSize = p.polyVecCompressedSize + polyCompressedSize
+
+	p.publicKeySize = p.indcpaPublicKeySize
+	p.secretKeySize = p.indcpaSecretKeySize + p.indcpaPublicKeySize + 2*SymSize // 32 bytes of additional space to save H(pk)
+	p.cipherTextSize = p.indcpaSize
+
+	return &p
+}

+ 146 - 0
poly.go

@@ -0,0 +1,146 @@
+// poly.go - Kyber polynomial.
+//
+// To the extent possible under law, Yawning Angel has waived all copyright
+// and related or neighboring rights to the software, using the Creative
+// Commons "CC0" public domain dedication. See LICENSE or
+// <http://creativecommons.org/publicdomain/zero/1.0/> for full details.
+
+package kyber
+
+import "golang.org/x/crypto/sha3"
+
+// Elements of R_q = Z_q[X]/(X^n + 1). Represents polynomial coeffs[0] +
+// X*coeffs[1] + X^2*xoeffs[2] + ... + X^{n-1}*coeffs[n-1].
+type poly struct {
+	coeffs [kyberN]uint16
+}
+
+// Compression and subsequent serialization of a polynomial.
+func (p *poly) compress(r []byte) {
+	var t [8]uint32
+
+	for i, k := 0, 0; i < kyberN; i, k = i+8, k+3 {
+		for j := 0; j < 8; j++ {
+			t[j] = uint32((((freeze(p.coeffs[i+j]) << 3) + kyberQ/2) / kyberQ) & 7)
+		}
+
+		r[k] = byte(t[0] | (t[1] << 3) | (t[2] << 6))
+		r[k+1] = byte((t[2] >> 2) | (t[3] << 1) | (t[4] << 4) | (t[5] << 7))
+		r[k+2] = byte((t[5] >> 1) | (t[6] << 2) | (t[7] << 5))
+	}
+}
+
+// De-serialization and subsequent decompression of a polynomial; approximate
+// inverse of poly.compress().
+func (p *poly) decompress(a []byte) {
+	for i, off := 0, 0; i < kyberN; i, off = i+8, off+3 {
+		p.coeffs[i+0] = ((uint16(a[off]&7) * kyberQ) + 4) >> 3
+		p.coeffs[i+1] = (((uint16(a[off]>>3) & 7) * kyberQ) + 4) >> 3
+		p.coeffs[i+2] = (((uint16(a[off]>>6) | (uint16(a[off+1]<<2) & 4)) * kyberQ) + 4) >> 3
+		p.coeffs[i+3] = (((uint16(a[off+1]>>1) & 7) * kyberQ) + 4) >> 3
+		p.coeffs[i+4] = (((uint16(a[off+1]>>4) & 7) * kyberQ) + 4) >> 3
+		p.coeffs[i+5] = (((uint16(a[off+1]>>7) | (uint16(a[off+2]<<1) & 6)) * kyberQ) + 4) >> 3
+		p.coeffs[i+6] = (((uint16(a[off+2]>>2) & 7) * kyberQ) + 4) >> 3
+		p.coeffs[i+7] = (((uint16(a[off+2] >> 5)) * kyberQ) + 4) >> 3
+	}
+}
+
+// Serialization of a polynomial.
+func (p *poly) toBytes(r []byte) {
+	var t [8]uint16
+
+	for i := 0; i < kyberN/8; i++ {
+		for j := 0; j < 8; j++ {
+			t[j] = freeze(p.coeffs[8*i+j])
+		}
+
+		r[13*i+0] = byte(t[0] & 0xff)
+		r[13*i+1] = byte((t[0] >> 8) | ((t[1] & 0x07) << 5))
+		r[13*i+2] = byte((t[1] >> 3) & 0xff)
+		r[13*i+3] = byte((t[1] >> 11) | ((t[2] & 0x3f) << 2))
+		r[13*i+4] = byte((t[2] >> 6) | ((t[3] & 0x01) << 7))
+		r[13*i+5] = byte((t[3] >> 1) & 0xff)
+		r[13*i+6] = byte((t[3] >> 9) | ((t[4] & 0x0f) << 4))
+		r[13*i+7] = byte((t[4] >> 4) & 0xff)
+		r[13*i+8] = byte((t[4] >> 12) | ((t[5] & 0x7f) << 1))
+		r[13*i+9] = byte((t[5] >> 7) | ((t[6] & 0x03) << 6))
+		r[13*i+10] = byte((t[6] >> 2) & 0xff)
+		r[13*i+11] = byte((t[6] >> 10) | ((t[7] & 0x1f) << 3))
+		r[13*i+12] = byte(t[7] >> 5)
+	}
+}
+
+// De-serialization of a polynomial; inverse of poly.toBytes().
+func (p *poly) fromBytes(a []byte) {
+	for i := 0; i < kyberN/8; i++ {
+		p.coeffs[8*i+0] = uint16(a[13*i+0]) | ((uint16(a[13*i+1]) & 0x1f) << 8)
+		p.coeffs[8*i+1] = (uint16(a[13*i+1]) >> 5) | (uint16(a[13*i+2]) << 3) | ((uint16(a[13*i+3]) & 0x03) << 11)
+		p.coeffs[8*i+2] = (uint16(a[13*i+3]) >> 2) | ((uint16(a[13*i+4]) & 0x7f) << 6)
+		p.coeffs[8*i+3] = (uint16(a[13*i+4]) >> 7) | (uint16(a[13*i+5]) << 1) | ((uint16(a[13*i+6]) & 0x0f) << 9)
+		p.coeffs[8*i+4] = (uint16(a[13*i+6]) >> 4) | (uint16(a[13*i+7]) << 4) | ((uint16(a[13*i+8]) & 0x01) << 12)
+		p.coeffs[8*i+5] = (uint16(a[13*i+8]) >> 1) | ((uint16(a[13*i+9]) & 0x3f) << 7)
+		p.coeffs[8*i+6] = (uint16(a[13*i+9]) >> 6) | (uint16(a[13*i+10]) << 2) | ((uint16(a[13*i+11]) & 0x07) << 10)
+		p.coeffs[8*i+7] = (uint16(a[13*i+11]) >> 3) | (uint16(a[13*i+12]) << 5)
+	}
+}
+
+// Convert 32-byte message to polynomial.
+func (p *poly) fromMsg(msg []byte) {
+	for i, v := range msg[:SymSize] {
+		for j := 0; j < 8; j++ {
+			mask := -((uint16(v) >> uint(j)) & 1)
+			p.coeffs[8*i+j] = mask & ((kyberQ + 1) / 2)
+		}
+	}
+}
+
+// Convert polynomial to 32-byte message.
+func (p *poly) toMsg(msg []byte) {
+	for i := 0; i < SymSize; i++ {
+		msg[i] = 0
+		for j := 0; j < 8; j++ {
+			t := (((freeze(p.coeffs[8*i+j]) << 1) + kyberQ/2) / kyberQ) & 1
+			msg[i] |= byte(t << uint(j))
+		}
+	}
+}
+
+// Sample a polynomial deterministically from a seed and a nonce, with output
+// polynomial close to centered binomial distribution with parameter eta.
+func (p *poly) getNoise(seed []byte, nonce byte, eta int) {
+	extSeed := make([]byte, 0, SymSize+1)
+	extSeed = append(extSeed, seed...)
+	extSeed = append(extSeed, nonce)
+
+	buf := make([]byte, eta*kyberN/4)
+	sha3.ShakeSum256(buf, extSeed)
+
+	p.cbd(buf, eta)
+}
+
+// Computes negacyclic number-theoretic transform (NTT) of a polynomial in
+// place; inputs assumed to be in normal order, output in bitreversed order.
+func (p *poly) ntt() {
+	ntt(&p.coeffs)
+}
+
+// Computes inverse of negacyclic number-theoretic transform (NTT) of a
+// polynomial in place; inputs assumed to be in bitreversed order, output in
+// normal order.
+func (p *poly) invntt() {
+	invntt(&p.coeffs)
+}
+
+// Add two polynomials.
+func (p *poly) add(a, b *poly) {
+	for i := range p.coeffs {
+		p.coeffs[i] = barrettReduce(a.coeffs[i] + b.coeffs[i])
+	}
+}
+
+// Subtract two polynomials.
+func (p *poly) sub(a, b *poly) {
+	for i := range p.coeffs {
+		p.coeffs[i] = barrettReduce(a.coeffs[i] + 3*kyberQ - b.coeffs[i])
+	}
+}

+ 110 - 0
polyvec.go

@@ -0,0 +1,110 @@
+// polyvec.go - Vector of Kyber polynomials.
+//
+// To the extent possible under law, Yawning Angel has waived all copyright
+// and related or neighboring rights to the software, using the Creative
+// Commons "CC0" public domain dedication. See LICENSE or
+// <http://creativecommons.org/publicdomain/zero/1.0/> for full details.
+
+package kyber
+
+type polyVec struct {
+	vec []*poly
+}
+
+// Compress and serialize vector of polynomials.
+func (v *polyVec) compress(r []byte) {
+	var off int
+	for _, vec := range v.vec {
+		for j := 0; j < kyberN/8; j++ {
+			var t [8]uint16
+			for k := 0; k < 8; k++ {
+				t[k] = uint16((((uint32(freeze(vec.coeffs[8*j+k])) << 11) + kyberQ/2) / kyberQ) & 0x7ff)
+			}
+
+			r[off+11*j+0] = byte(t[0] & 0xff)
+			r[off+11*j+1] = byte((t[0] >> 8) | ((t[1] & 0x1f) << 3))
+			r[off+11*j+2] = byte((t[1] >> 5) | ((t[2] & 0x03) << 6))
+			r[off+11*j+3] = byte((t[2] >> 2) & 0xff)
+			r[off+11*j+4] = byte((t[2] >> 10) | ((t[3] & 0x7f) << 1))
+			r[off+11*j+5] = byte((t[3] >> 7) | ((t[4] & 0x0f) << 4))
+			r[off+11*j+6] = byte((t[4] >> 4) | ((t[5] & 0x01) << 7))
+			r[off+11*j+7] = byte((t[5] >> 1) & 0xff)
+			r[off+11*j+8] = byte((t[5] >> 9) | ((t[6] & 0x3f) << 2))
+			r[off+11*j+9] = byte((t[6] >> 6) | ((t[7] & 0x07) << 5))
+			r[off+11*j+10] = byte((t[7] >> 3))
+		}
+		off += compressedCoeffSize
+	}
+}
+
+// De-serialize and decompress vector of polynomials; approximate inverse of
+// polyVec.compress().
+func (v *polyVec) decompress(a []byte) {
+	var off int
+	for _, vec := range v.vec {
+		for j := 0; j < kyberN/8; j++ {
+			vec.coeffs[8*j+0] = uint16((((uint32(a[off+11*j+0]) | ((uint32(a[off+11*j+1]) & 0x07) << 8)) * kyberQ) + 1024) >> 11)
+			vec.coeffs[8*j+1] = uint16(((((uint32(a[off+11*j+1]) >> 3) | ((uint32(a[off+11*j+2]) & 0x3f) << 5)) * kyberQ) + 1024) >> 11)
+			vec.coeffs[8*j+2] = uint16(((((uint32(a[off+11*j+2]) >> 6) | ((uint32(a[off+11*j+3]) & 0xff) << 2) | ((uint32(a[off+11*j+4]) & 0x01) << 10)) * kyberQ) + 1024) >> 11)
+			vec.coeffs[8*j+3] = uint16(((((uint32(a[off+11*j+4]) >> 1) | ((uint32(a[off+11*j+5]) & 0x0f) << 7)) * kyberQ) + 1024) >> 11)
+			vec.coeffs[8*j+4] = uint16(((((uint32(a[off+11*j+5]) >> 4) | ((uint32(a[off+11*j+6]) & 0x7f) << 4)) * kyberQ) + 1024) >> 11)
+			vec.coeffs[8*j+5] = uint16(((((uint32(a[off+11*j+6]) >> 7) | ((uint32(a[off+11*j+7]) & 0xff) << 1) | ((uint32(a[off+11*j+8]) & 0x03) << 9)) * kyberQ) + 1024) >> 11)
+			vec.coeffs[8*j+6] = uint16(((((uint32(a[off+11*j+8]) >> 2) | ((uint32(a[off+11*j+9]) & 0x1f) << 6)) * kyberQ) + 1024) >> 11)
+			vec.coeffs[8*j+7] = uint16(((((uint32(a[off+11*j+9]) >> 5) | ((uint32(a[off+11*j+10]) & 0xff) << 3)) * kyberQ) + 1024) >> 11)
+		}
+		off += compressedCoeffSize
+	}
+}
+
+// Serialize vector of polynomials.
+func (v *polyVec) toBytes(r []byte) {
+	for i, p := range v.vec {
+		p.toBytes(r[i*polySize:])
+	}
+}
+
+// De-serialize vector of polynomials; inverse of polyVec.toBytes().
+func (v *polyVec) fromBytes(a []byte) {
+	for i, p := range v.vec {
+		p.fromBytes(a[i*polySize:])
+	}
+}
+
+// Apply forward NTT to all elements of a vector of polynomials.
+func (v *polyVec) ntt() {
+	for _, p := range v.vec {
+		p.ntt()
+	}
+}
+
+// Apply inverse NTT to all elements of a vector of polynomials.
+func (v *polyVec) invntt() {
+	for _, p := range v.vec {
+		p.invntt()
+	}
+}
+
+// Pointwise multiply elements of a and b and accumulate into p.
+func (p *poly) pointwiseAcc(a, b *polyVec) {
+	for j := 0; j < kyberN; j++ {
+		t := montgomeryReduce(4613 * uint32(b.vec[0].coeffs[j])) // 4613 = 2^{2*18} % q
+		p.coeffs[j] = montgomeryReduce(uint32(a.vec[0].coeffs[j]) * uint32(t))
+		for i := 1; i < len(a.vec); i++ { // len(a.vec) == kyberK
+			t = montgomeryReduce(4613 * uint32(b.vec[i].coeffs[j]))
+			p.coeffs[j] += montgomeryReduce(uint32(a.vec[i].coeffs[j]) * uint32(t))
+		}
+		p.coeffs[j] = barrettReduce(p.coeffs[j])
+	}
+}
+
+// Add vectors of polynomials.
+func (v *polyVec) add(a, b *polyVec) {
+	for i, p := range v.vec {
+		p.add(a.vec[i], b.vec[i])
+	}
+}
+
+// Get compressed and serialized size in bytes.
+func (v *polyVec) compressedSize() int {
+	return len(v.vec) * compressedCoeffSize
+}

+ 87 - 0
precomp.go

@@ -0,0 +1,87 @@
+// precomp.go - Precomputed NTT constants.
+//
+// To the extent possible under law, Yawning Angel has waived all copyright
+// and related or neighboring rights to the software, using the Creative
+// Commons "CC0" public domain dedication. See LICENSE or
+// <http://creativecommons.org/publicdomain/zero/1.0/> for full details.
+
+package kyber
+
+// Precomputed constants for the forward NTT and inverse NTT.
+// Computed using Pari/GP as follows:
+//
+// brv=[0,128,64,192,32,160,96,224,16,144,80,208,48,176,112,240, \
+//      8,136,72,200,40,168,104,232,24,152,88,216,56,184,120,248, \
+//      4,132,68,196,36,164,100,228,20,148,84,212,52,180,116,244, \
+//      12,140,76,204,44,172,108,236,28,156,92,220,60,188,124,252, \
+//      2,130,66,194,34,162,98,226,18,146,82,210,50,178,114,242, \
+//      10,138,74,202,42,170,106,234,26,154,90,218,58,186,122,250, \
+//      6,134,70,198,38,166,102,230,22,150,86,214,54,182,118,246, \
+//      14,142,78,206,46,174,110,238,30,158,94,222,62,190,126,254, \
+//      1,129,65,193,33,161,97,225,17,145,81,209,49,177,113,241, \
+//      9,137,73,201,41,169,105,233,25,153,89,217,57,185,121,249, \
+//      5,133,69,197,37,165,101,229,21,149,85,213,53,181,117,245, \
+//      13,141,77,205,45,173,109,237,29,157,93,221,61,189,125,253, \
+//      3,131,67,195,35,163,99,227,19,147,83,211,51,179,115,243, \
+//      11,139,75,203,43,171,107,235,27,155,91,219,59,187,123,251, \
+//      7,135,71,199,39,167,103,231,23,151,87,215,55,183,119,247, \
+//      15,143,79,207,47,175,111,239,31,159,95,223,63,191,127,255];
+//
+// q = 7681;
+// n = 256;
+// mont = Mod(2^18,q);
+//
+// g=0; for(i=2,q-1,if(znorder(Mod(i,q)) == 2*n, g=Mod(i,q); break))
+//
+// zetas = lift(vector(n, i, g^(brv[i])*mont))
+// omegas_inv_bitrev_montgomery = lift(vector(n/2, i, (g^2)^(-brv[2*(i-1)+1])*mont))
+// psis_inv_montgomery = lift(vector(n, i, g^(-(i-1))/n*mont))
+
+var zetas = [kyberN]uint16{
+	990, 7427, 2634, 6819, 578, 3281, 2143, 1095, 484, 6362, 3336, 5382, 6086, 3823, 877, 5656,
+	3583, 7010, 6414, 263, 1285, 291, 7143, 7338, 1581, 5134, 5184, 5932, 4042, 5775, 2468, 3,
+	606, 729, 5383, 962, 3240, 7548, 5129, 7653, 5929, 4965, 2461, 641, 1584, 2666, 1142, 157,
+	7407, 5222, 5602, 5142, 6140, 5485, 4931, 1559, 2085, 5284, 2056, 3538, 7269, 3535, 7190, 1957,
+	3465, 6792, 1538, 4664, 2023, 7643, 3660, 7673, 1694, 6905, 3995, 3475, 5939, 1859, 6910, 4434,
+	1019, 1492, 7087, 4761, 657, 4859, 5798, 2640, 1693, 2607, 2782, 5400, 6466, 1010, 957, 3851,
+	2121, 6392, 7319, 3367, 3659, 3375, 6430, 7583, 1549, 5856, 4773, 6084, 5544, 1650, 3997, 4390,
+	6722, 2915, 4245, 2635, 6128, 7676, 5737, 1616, 3457, 3132, 7196, 4702, 6239, 851, 2122, 3009,
+	7613, 7295, 2007, 323, 5112, 3716, 2289, 6442, 6965, 2713, 7126, 3401, 963, 6596, 607, 5027,
+	7078, 4484, 5937, 944, 2860, 2680, 5049, 1777, 5850, 3387, 6487, 6777, 4812, 4724, 7077, 186,
+	6848, 6793, 3463, 5877, 1174, 7116, 3077, 5945, 6591, 590, 6643, 1337, 6036, 3991, 1675, 2053,
+	6055, 1162, 1679, 3883, 4311, 2106, 6163, 4486, 6374, 5006, 4576, 4288, 5180, 4102, 282, 6119,
+	7443, 6330, 3184, 4971, 2530, 5325, 4171, 7185, 5175, 5655, 1898, 382, 7211, 43, 5965, 6073,
+	1730, 332, 1577, 3304, 2329, 1699, 6150, 2379, 5113, 333, 3502, 4517, 1480, 1172, 5567, 651,
+	925, 4573, 599, 1367, 4109, 1863, 6929, 1605, 3866, 2065, 4048, 839, 5764, 2447, 2022, 3345,
+	1990, 4067, 2036, 2069, 3567, 7371, 2368, 339, 6947, 2159, 654, 7327, 2768, 6676, 987, 2214,
+}
+
+var omegasInvBitrevMontgomery = [kyberN / 2]uint16{
+	990, 254, 862, 5047, 6586, 5538, 4400, 7103, 2025, 6804, 3858, 1595, 2299, 4345, 1319, 7197,
+	7678, 5213, 1906, 3639, 1749, 2497, 2547, 6100, 343, 538, 7390, 6396, 7418, 1267, 671, 4098,
+	5724, 491, 4146, 412, 4143, 5625, 2397, 5596, 6122, 2750, 2196, 1541, 2539, 2079, 2459, 274,
+	7524, 6539, 5015, 6097, 7040, 5220, 2716, 1752, 28, 2552, 133, 4441, 6719, 2298, 6952, 7075,
+	4672, 5559, 6830, 1442, 2979, 485, 4549, 4224, 6065, 1944, 5, 1553, 5046, 3436, 4766, 959,
+	3291, 3684, 6031, 2137, 1597, 2908, 1825, 6132, 98, 1251, 4306, 4022, 4314, 362, 1289, 5560,
+	3830, 6724, 6671, 1215, 2281, 4899, 5074, 5988, 5041, 1883, 2822, 7024, 2920, 594, 6189, 6662,
+	3247, 771, 5822, 1742, 4206, 3686, 776, 5987, 8, 4021, 38, 5658, 3017, 6143, 889, 4216,
+}
+
+var psisInvMontgomery = [kyberN]uint16{
+	1024, 4972, 5779, 6907, 4943, 4168, 315, 5580, 90, 497, 1123, 142, 4710, 5527, 2443, 4871,
+	698, 2489, 2394, 4003, 684, 2241, 2390, 7224, 5072, 2064, 4741, 1687, 6841, 482, 7441, 1235,
+	2126, 4742, 2802, 5744, 6287, 4933, 699, 3604, 1297, 2127, 5857, 1705, 3868, 3779, 4397, 2177,
+	159, 622, 2240, 1275, 640, 6948, 4572, 5277, 209, 2605, 1157, 7328, 5817, 3191, 1662, 2009,
+	4864, 574, 2487, 164, 6197, 4436, 7257, 3462, 4268, 4281, 3414, 4515, 3170, 1290, 2003, 5855,
+	7156, 6062, 7531, 1732, 3249, 4884, 7512, 3590, 1049, 2123, 1397, 6093, 3691, 6130, 6541, 3946,
+	6258, 3322, 1788, 4241, 4900, 2309, 1400, 1757, 400, 502, 6698, 2338, 3011, 668, 7444, 4580,
+	6516, 6795, 2959, 4136, 3040, 2279, 6355, 3943, 2913, 6613, 7416, 4084, 6508, 5556, 4054, 3782,
+	61, 6567, 2212, 779, 632, 5709, 5667, 4923, 4911, 6893, 4695, 4164, 3536, 2287, 7594, 2848,
+	3267, 1911, 3128, 546, 1991, 156, 4958, 5531, 6903, 483, 875, 138, 250, 2234, 2266, 7222,
+	2842, 4258, 812, 6703, 232, 5207, 6650, 2585, 1900, 6225, 4932, 7265, 4701, 3173, 4635, 6393,
+	227, 7313, 4454, 4284, 6759, 1224, 5223, 1447, 395, 2608, 4502, 4037, 189, 3348, 54, 6443,
+	2210, 6230, 2826, 1780, 3002, 5995, 1955, 6102, 6045, 3938, 5019, 4417, 1434, 1262, 1507, 5847,
+	5917, 7157, 7177, 6434, 7537, 741, 4348, 1309, 145, 374, 2236, 4496, 5028, 6771, 6923, 7421,
+	1978, 1023, 3857, 6876, 1102, 7451, 4704, 6518, 1344, 765, 384, 5705, 1207, 1630, 4734, 1563,
+	6839, 5933, 1954, 4987, 7142, 5814, 7527, 4953, 7637, 4707, 2182, 5734, 2818, 541, 4097, 5641,
+}

+ 43 - 0
reduce.go

@@ -0,0 +1,43 @@
+// reduce.go - Montgomery, Barret, and Full reduction.
+//
+// To the extent possible under law, Yawning Angel has waived all copyright
+// and related or neighboring rights to the software, using the Creative
+// Commons "CC0" public domain dedication. See LICENSE or
+// <http://creativecommons.org/publicdomain/zero/1.0/> for full details.
+
+package kyber
+
+const (
+	qinv = 7679 // -inverse_mod(q,2^18)
+	rlog = 18
+)
+
+// Montgomery reduction; given a 32-bit integer a, computes 16-bit integer
+// congruent to a * R^-1 mod q, where R=2^18 (see value of rlog).
+func montgomeryReduce(a uint32) uint16 {
+	u := a * qinv
+	u &= (1 << rlog) - 1
+	u *= kyberQ
+	a += u
+	return uint16(a >> rlog)
+}
+
+// Barrett reduction; given a 16-bit integer a, computes 16-bit integer
+// congruent to a mod q in {0,...,11768}.
+func barrettReduce(a uint16) uint16 {
+	u := uint32(a >> 13) // ((uint32_t) a * sinv) >> 16
+	u *= kyberQ
+	a -= uint16(u)
+	return a
+}
+
+// Full reduction; given a 16-bit integer a, computes unsigned integer a mod q.
+func freeze(x uint16) uint16 {
+	r := barrettReduce(x)
+
+	m := r - kyberQ
+	c := int16(m)
+	c >>= 15
+	r = m ^ ((r ^ m) & uint16(c))
+	return r
+}

+ 1 - 0
testdata/.gitignore

@@ -0,0 +1 @@
+*.full

+ 17 - 0
testdata/README.testdata

@@ -0,0 +1,17 @@
+The full test vectors generated by the reference code are rather large
+(approximately 28 MiB) and can not be comfortably included in the git
+repository.
+
+Instead, a compact representation containing the SHA256 digest of each
+test vector set (compactVectors.json) will be used by default.
+
+The tests will load and use the full test vector output if available.
+The files can be generated from the reference implementation by something
+like:
+
+  $ ./testvectors512 > KEM-Kyber-512.full
+  $ ./testvectors768 > KEM-Kyber-768.full
+  $ ./testvectors1024 > KEM-Kyber-1024.full
+
+  [Copy the `.full` files to `testdata/`.]
+

+ 5 - 0
testdata/compactVectors.json

@@ -0,0 +1,5 @@
+{
+  "Kyber-512": "6973360b86dbd7ceaef621e913cba98f2e95cd335a7cf22cca0f24dbb7e47da4",
+  "Kyber-768": "8f1673ccf9db0851dbe4826ade672e54441a59969994814eb35bd231628f28ee",
+  "Kyber-1024": "1c3839c0cbbe65c6decf443720efa9b949b9eef5a537b4cc421f22334c91c06a"
+}