crypto/internal/fips/mlkem: implement ML-KEM-1024

Decided to automatically duplicate the high-level code to avoid growing
the ML-KEM-768 data structures.

For #70122

Change-Id: I5c705b71ee1e23adba9113d5cf6b6e505c028967
Reviewed-on: https://go-review.googlesource.com/c/go/+/621983
Auto-Submit: Filippo Valsorda <filippo@golang.org>
Reviewed-by: Roland Shoemaker <roland@golang.org>
Reviewed-by: Daniel McCarney <daniel@binaryparadox.net>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Dmitri Shuralyov <dmitshur@google.com>

Author: FiloSottile

src/crypto/internal/fips/mlkem/field.go

@@ -263,7 +263,7 @@ func ringCompressAndEncode10(s []byte, f ringElement) []byte {
 	s, b := sliceForAppend(s, encodingSize10)
 	for i := 0; i < n; i += 4 {
 		var x uint64
-		x |= uint64(compress(f[i+0], 10))
+		x |= uint64(compress(f[i], 10))
 		x |= uint64(compress(f[i+1], 10)) << 10
 		x |= uint64(compress(f[i+2], 10)) << 20
 		x |= uint64(compress(f[i+3], 10)) << 30
@@ -296,6 +296,101 @@ func ringDecodeAndDecompress10(bb *[encodingSize10]byte) ringElement {
 	return f
 }
 
+// ringCompressAndEncode appends an encoding of a ring element to s,
+// compressing each coefficient to d bits.
+//
+// It implements Compress, according to FIPS 203, Definition 4.7,
+// followed by ByteEncode, according to FIPS 203, Algorithm 5.
+func ringCompressAndEncode(s []byte, f ringElement, d uint8) []byte {
+	var b byte
+	var bIdx uint8
+	for i := 0; i < n; i++ {
+		c := compress(f[i], d)
+		var cIdx uint8
+		for cIdx < d {
+			b |= byte(c>>cIdx) << bIdx
+			bits := min(8-bIdx, d-cIdx)
+			bIdx += bits
+			cIdx += bits
+			if bIdx == 8 {
+				s = append(s, b)
+				b = 0
+				bIdx = 0
+			}
+		}
+	}
+	if bIdx != 0 {
+		panic("mlkem: internal error: bitsFilled != 0")
+	}
+	return s
+}
+
+// ringDecodeAndDecompress decodes an encoding of a ring element where
+// each d bits are mapped to an equidistant distribution.
+//
+// It implements ByteDecode, according to FIPS 203, Algorithm 6,
+// followed by Decompress, according to FIPS 203, Definition 4.8.
+func ringDecodeAndDecompress(b []byte, d uint8) ringElement {
+	var f ringElement
+	var bIdx uint8
+	for i := 0; i < n; i++ {
+		var c uint16
+		var cIdx uint8
+		for cIdx < d {
+			c |= uint16(b[0]>>bIdx) << cIdx
+			c &= (1 << d) - 1
+			bits := min(8-bIdx, d-cIdx)
+			bIdx += bits
+			cIdx += bits
+			if bIdx == 8 {
+				b = b[1:]
+				bIdx = 0
+			}
+		}
+		f[i] = fieldElement(decompress(c, d))
+	}
+	if len(b) != 0 {
+		panic("mlkem: internal error: leftover bytes")
+	}
+	return f
+}
+
+// ringCompressAndEncode5 appends a 160-byte encoding of a ring element to s,
+// compressing eight coefficients per five bytes.
+//
+// It implements Compress₅, according to FIPS 203, Definition 4.7,
+// followed by ByteEncode₅, according to FIPS 203, Algorithm 5.
+func ringCompressAndEncode5(s []byte, f ringElement) []byte {
+	return ringCompressAndEncode(s, f, 5)
+}
+
+// ringDecodeAndDecompress5 decodes a 160-byte encoding of a ring element where
+// each five bits are mapped to an equidistant distribution.
+//
+// It implements ByteDecode₅, according to FIPS 203, Algorithm 6,
+// followed by Decompress₅, according to FIPS 203, Definition 4.8.
+func ringDecodeAndDecompress5(bb *[encodingSize5]byte) ringElement {
+	return ringDecodeAndDecompress(bb[:], 5)
+}
+
+// ringCompressAndEncode11 appends a 352-byte encoding of a ring element to s,
+// compressing eight coefficients per eleven bytes.
+//
+// It implements Compress₁₁, according to FIPS 203, Definition 4.7,
+// followed by ByteEncode₁₁, according to FIPS 203, Algorithm 5.
+func ringCompressAndEncode11(s []byte, f ringElement) []byte {
+	return ringCompressAndEncode(s, f, 11)
+}
+
+// ringDecodeAndDecompress11 decodes a 352-byte encoding of a ring element where
+// each eleven bits are mapped to an equidistant distribution.
+//
+// It implements ByteDecode₁₁, according to FIPS 203, Algorithm 6,
+// followed by Decompress₁₁, according to FIPS 203, Definition 4.8.
+func ringDecodeAndDecompress11(bb *[encodingSize11]byte) ringElement {
+	return ringDecodeAndDecompress(bb[:], 11)
+}
+
 // samplePolyCBD draws a ringElement from the special Dη distribution given a
 // stream of random bytes generated by the PRF function, according to FIPS 203,
 // Algorithm 8 and Definition 4.3.

src/crypto/internal/fips/mlkem/field_test.go

@@ -5,7 +5,10 @@
 package mlkem
 
 import (
+	"bytes"
+	"crypto/rand"
 	"math/big"
+	mathrand "math/rand/v2"
 	"strconv"
 	"testing"
 )
@@ -151,6 +154,81 @@ func TestDecompress(t *testing.T) {
 	}
 }
 
+func randomRingElement() ringElement {
+	var r ringElement
+	for i := range r {
+		r[i] = fieldElement(mathrand.IntN(q))
+	}
+	return r
+}
+
+func TestEncodeDecode(t *testing.T) {
+	f := randomRingElement()
+	b := make([]byte, 12*n/8)
+	rand.Read(b)
+
+	// Compare ringCompressAndEncode to ringCompressAndEncodeN.
+	e1 := ringCompressAndEncode(nil, f, 10)
+	e2 := ringCompressAndEncode10(nil, f)
+	if !bytes.Equal(e1, e2) {
+		t.Errorf("ringCompressAndEncode = %x, ringCompressAndEncode10 = %x", e1, e2)
+	}
+	e1 = ringCompressAndEncode(nil, f, 4)
+	e2 = ringCompressAndEncode4(nil, f)
+	if !bytes.Equal(e1, e2) {
+		t.Errorf("ringCompressAndEncode = %x, ringCompressAndEncode4 = %x", e1, e2)
+	}
+	e1 = ringCompressAndEncode(nil, f, 1)
+	e2 = ringCompressAndEncode1(nil, f)
+	if !bytes.Equal(e1, e2) {
+		t.Errorf("ringCompressAndEncode = %x, ringCompressAndEncode1 = %x", e1, e2)
+	}
+
+	// Compare ringDecodeAndDecompress to ringDecodeAndDecompressN.
+	g1 := ringDecodeAndDecompress(b[:encodingSize10], 10)
+	g2 := ringDecodeAndDecompress10((*[encodingSize10]byte)(b))
+	if g1 != g2 {
+		t.Errorf("ringDecodeAndDecompress = %v, ringDecodeAndDecompress10 = %v", g1, g2)
+	}
+	g1 = ringDecodeAndDecompress(b[:encodingSize4], 4)
+	g2 = ringDecodeAndDecompress4((*[encodingSize4]byte)(b))
+	if g1 != g2 {
+		t.Errorf("ringDecodeAndDecompress = %v, ringDecodeAndDecompress4 = %v", g1, g2)
+	}
+	g1 = ringDecodeAndDecompress(b[:encodingSize1], 1)
+	g2 = ringDecodeAndDecompress1((*[encodingSize1]byte)(b))
+	if g1 != g2 {
+		t.Errorf("ringDecodeAndDecompress = %v, ringDecodeAndDecompress1 = %v", g1, g2)
+	}
+
+	// Round-trip ringCompressAndEncode and ringDecodeAndDecompress.
+	for d := 1; d < 12; d++ {
+		encodingSize := d * n / 8
+		g := ringDecodeAndDecompress(b[:encodingSize], uint8(d))
+		out := ringCompressAndEncode(nil, g, uint8(d))
+		if !bytes.Equal(out, b[:encodingSize]) {
+			t.Errorf("roundtrip failed for d = %d", d)
+		}
+	}
+
+	// Round-trip ringCompressAndEncodeN and ringDecodeAndDecompressN.
+	g := ringDecodeAndDecompress10((*[encodingSize10]byte)(b))
+	out := ringCompressAndEncode10(nil, g)
+	if !bytes.Equal(out, b[:encodingSize10]) {
+		t.Errorf("roundtrip failed for specialized 10")
+	}
+	g = ringDecodeAndDecompress4((*[encodingSize4]byte)(b))
+	out = ringCompressAndEncode4(nil, g)
+	if !bytes.Equal(out, b[:encodingSize4]) {
+		t.Errorf("roundtrip failed for specialized 4")
+	}
+	g = ringDecodeAndDecompress1((*[encodingSize1]byte)(b))
+	out = ringCompressAndEncode1(nil, g)
+	if !bytes.Equal(out, b[:encodingSize1]) {
+		t.Errorf("roundtrip failed for specialized 1")
+	}
+}
+
 func BitRev7(n uint8) uint8 {
 	if n>>7 != 0 {
 		panic("not 7 bits")

src/crypto/internal/fips/mlkem/generate1024.go

@@ -0,0 +1,123 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build ignore
+
+package main
+
+import (
+	"flag"
+	"go/ast"
+	"go/format"
+	"go/parser"
+	"go/token"
+	"log"
+	"os"
+	"strings"
+)
+
+var replacements = map[string]string{
+	"k": "k1024",
+
+	"CiphertextSize768":       "CiphertextSize1024",
+	"EncapsulationKeySize768": "EncapsulationKeySize1024",
+
+	"encryptionKey": "encryptionKey1024",
+	"decryptionKey": "decryptionKey1024",
+
+	"EncapsulationKey768":    "EncapsulationKey1024",
+	"NewEncapsulationKey768": "NewEncapsulationKey1024",
+	"parseEK":                "parseEK1024",
+
+	"kemEncaps":  "kemEncaps1024",
+	"pkeEncrypt": "pkeEncrypt1024",
+
+	"DecapsulationKey768":    "DecapsulationKey1024",
+	"NewDecapsulationKey768": "NewDecapsulationKey1024",
+	"newKeyFromSeed":         "newKeyFromSeed1024",
+
+	"kemDecaps":  "kemDecaps1024",
+	"pkeDecrypt": "pkeDecrypt1024",
+
+	"GenerateKey768": "GenerateKey1024",
+	"generateKey":    "generateKey1024",
+
+	"kemKeyGen": "kemKeyGen1024",
+
+	"encodingSize4":             "encodingSize5",
+	"encodingSize10":            "encodingSize11",
+	"ringCompressAndEncode4":    "ringCompressAndEncode5",
+	"ringCompressAndEncode10":   "ringCompressAndEncode11",
+	"ringDecodeAndDecompress4":  "ringDecodeAndDecompress5",
+	"ringDecodeAndDecompress10": "ringDecodeAndDecompress11",
+}
+
+func main() {
+	inputFile := flag.String("input", "", "")
+	outputFile := flag.String("output", "", "")
+	flag.Parse()
+
+	fset := token.NewFileSet()
+	f, err := parser.ParseFile(fset, *inputFile, nil, parser.SkipObjectResolution|parser.ParseComments)
+	if err != nil {
+		log.Fatal(err)
+	}
+	cmap := ast.NewCommentMap(fset, f, f.Comments)
+
+	// Drop header comments.
+	cmap[ast.Node(f)] = nil
+
+	// Remove top-level consts used across the main and generated files.
+	var newDecls []ast.Decl
+	for _, decl := range f.Decls {
+		switch d := decl.(type) {
+		case *ast.GenDecl:
+			if d.Tok == token.CONST {
+				continue // Skip const declarations
+			}
+			if d.Tok == token.IMPORT {
+				cmap[decl] = nil // Drop pre-import comments.
+			}
+		}
+		newDecls = append(newDecls, decl)
+	}
+	f.Decls = newDecls
+
+	// Replace identifiers.
+	ast.Inspect(f, func(n ast.Node) bool {
+		switch x := n.(type) {
+		case *ast.Ident:
+			if replacement, ok := replacements[x.Name]; ok {
+				x.Name = replacement
+			}
+		}
+		return true
+	})
+
+	// Replace identifiers in comments.
+	for _, c := range f.Comments {
+		for _, l := range c.List {
+			for k, v := range replacements {
+				if k == "k" {
+					continue
+				}
+				l.Text = strings.ReplaceAll(l.Text, k, v)
+			}
+		}
+	}
+
+	out, err := os.Create(*outputFile)
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer out.Close()
+
+	out.WriteString("// Code generated by generate1024.go. DO NOT EDIT.\n\n")
+
+	f.Comments = cmap.Filter(f).Comments()
+	err = format.Node(out, fset, f)
+	if err != nil {
+		log.Fatal(err)
+	}
+}