use sha3-512 for drbg

README.md

@@ -7,7 +7,7 @@ It's possible there are a few timing vulnerabilities, and I've only just learned
 However, I did my best to implement as instructed.
 
 Future improvements:
-- [ ] Stop using Sha256 to provide entropy, use Sha3-256 or Blake3 or something approved and better
+- [x] Stop using Sha256 to provide entropy, use Sha3-256 or Blake3 or something approved and better (chose SHA3-512, I dunno I don't think it's necessary to use 512 but meh)
 - [ ] Use hardware for NTT math
   - [x] NEON (arm) update: partial implementation sees negative performance gain, but maybe I'm just doing it wrong. if you want to try building it, just build with `-tags=neon`. not planning on investing more time here.
   - [ ] AVX (x86)

entropy.go

@@ -1,19 +1,151 @@
 package mldsa
 
 import (
-	"crypto"
+	"crypto/rand"
+	"encoding/hex"
+	"fmt"
+	"math/big"
 
-	drbg "github.com/canonical/go-sp800.90a-drbg"
+	"golang.org/x/crypto/sha3"
 )
 
+type DRBG struct {
+	seedLength    int32
+	reseedCounter int64
+
+	V []byte
+	C []byte
+}
+
+func (drbg *DRBG) Init(entropy []byte, personalizationString []byte) {
+	drbg.seedLength = 888
+
+	seedMaterial := append(entropy, personalizationString...)
+	drbg.reseed(seedMaterial)
+}
+
+func (drbg *DRBG) reseed(seedMaterial []byte) {
+	drbg.V = drbg.derive(seedMaterial, drbg.seedLength)
+	drbg.C = drbg.derive(append([]byte{1}, drbg.V...), drbg.seedLength)
+	drbg.reseedCounter += 1
+}
+
+func (drbg *DRBG) derive(input []byte, numberOfBits int32) []byte {
+	numberOfBytes := (numberOfBits + 7) / 8
+
+	output := []byte{}
+	temp := input[:]
+
+	offset := int32(0)
+
+	for offset < numberOfBytes {
+		temp = append([]byte{1}, temp...)
+		dig := sha3.Sum512(temp)
+		digLength := int32(64)
+
+		var bytesToWrite int32
+		if digLength < numberOfBytes-offset {
+			bytesToWrite = digLength
+		} else {
+			bytesToWrite = numberOfBytes - offset
+		}
+
+		output = append(output, dig[:bytesToWrite]...)
+		offset += bytesToWrite
+	}
+
+	return output
+}
+
+func (drbg *DRBG) Reseed(entropy []byte) {
+	seedMaterial := append([]byte{1}, drbg.V...)
+	seedMaterial = append(seedMaterial, entropy...)
+	drbg.reseed(seedMaterial)
+}
+
+func (drbg *DRBG) Generate(numberOfBits int32) ([]byte, error) {
+	if drbg.reseedCounter > int64(1<<48) {
+		return nil, fmt.Errorf("must reseed")
+	}
+
+	numberOfBytes := (numberOfBits + 7) / 8
+	output := []byte{}
+
+	offset := int32(0)
+	for offset < numberOfBytes {
+		sum512 := sha3.Sum512(drbg.V)
+		drbg.V = sum512[:]
+		digLength := int32(64)
+
+		var bytesToWrite int32
+		if digLength < numberOfBytes-offset {
+			bytesToWrite = digLength
+		} else {
+			bytesToWrite = numberOfBytes - offset
+		}
+
+		output = append(output, drbg.V[:bytesToWrite]...)
+		offset += bytesToWrite
+	}
+
+	V := bytesToBigInt(drbg.V)
+	C := bytesToBigInt(drbg.C)
+	sum := &big.Int{}
+
+	sum.Add(V, C)
+	X := *sum
+	V = &X
+	sum.Add(V, big.NewInt(int64(drbg.reseedCounter)))
+	drbg.V = bigIntToBytes(sum, int(drbg.seedLength/8))
+
+	drbg.reseedCounter += 1
+
+	return output, nil
+}
+
+func bytesToBigInt(bytes []byte) *big.Int {
+	hexStr := hex.EncodeToString(bytes)
+	bigInt := new(big.Int)
+	bigInt.SetString(hexStr, 16)
+	return bigInt
+}
+
+func bigIntToBytes(number *big.Int, length int) []byte {
+	hexStr := number.Text(16)
+	hexStr = padLeft(hexStr, length*2, '0')
+
+	bytes, _ := hex.DecodeString(hexStr)
+
+	if len(bytes) < length {
+		paddedBytes := make([]byte, length)
+		copy(paddedBytes[length-len(bytes):], bytes)
+		return paddedBytes
+	}
+	return bytes[:length]
+}
+
+func padLeft(str string, length int, pad byte) string {
+	if len(str) >= length {
+		return str
+	}
+	padding := make([]byte, length-len(str))
+	for i := range padding {
+		padding[i] = pad
+	}
+	return string(padding) + str
+}
+
 func rbg(len int32) ([]byte, error) {
-	hashRbg, err := drbg.NewHash(crypto.SHA256, []byte{}, nil)
+	entropy := make([]byte, 32)
+	_, err := rand.Read(entropy)
 	if err != nil {
 		return nil, err
 	}
 
-	bytes := make([]byte, len)
-	err = hashRbg.Generate([]byte{}, bytes)
+	hashDRBG := &DRBG{}
+	hashDRBG.Init(entropy, []byte("jasoncolburne/ml-dsa-go"))
+
+	bytes, err := hashDRBG.Generate(len * 8)
 	if err != nil {
 		return nil, err
 	}

go.mod

@@ -2,9 +2,6 @@ module github.com/jasoncolburne/ml-dsa-go
 
 go 1.23.4
 
-require (
-	github.com/canonical/go-sp800.90a-drbg v0.0.0-20210314144037-6eeb1040d6c3 // indirect
-	golang.org/x/crypto v0.31.0 // indirect
-	golang.org/x/sys v0.28.0 // indirect
-	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect
-)
+require golang.org/x/crypto v0.31.0
+
+require golang.org/x/sys v0.28.0 // indirect

go.sum

@@ -1,12 +1,4 @@
-github.com/canonical/go-sp800.90a-drbg v0.0.0-20210314144037-6eeb1040d6c3 h1:oe6fCvaEpkhyW3qAicT0TnGtyht/UrgvOwMcEgLb7Aw=
-github.com/canonical/go-sp800.90a-drbg v0.0.0-20210314144037-6eeb1040d6c3/go.mod h1:qdP0gaj0QtgX2RUZhnlVrceJ+Qln8aSlDyJwelLLFeM=
-github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
-github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
-github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 golang.org/x/crypto v0.31.0 h1:ihbySMvVjLAeSH1IbfcRTkD/iNscyz8rGzjF/E5hV6U=
 golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk=
 golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA=
 golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
-golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE=
-golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=