From 86a19d04b7f54178b5520d1ca4fe6e73a3d35516 Mon Sep 17 00:00:00 2001
From: Till Faelligen <2353100+S7evinK@users.noreply.github.com>
Date: Mon, 22 Aug 2022 16:39:27 +0200
Subject: [PATCH] Add option to specify keySize when generating keys

---
 cmd/generate-keys/main.go |  5 +++--
 test/http.go              |  2 +-
 test/keys.go              | 12 ++++++------
 3 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/cmd/generate-keys/main.go b/cmd/generate-keys/main.go
index 8acd28be0..d4c8cf78a 100644
--- a/cmd/generate-keys/main.go
+++ b/cmd/generate-keys/main.go
@@ -38,6 +38,7 @@ var (
 	authorityCertFile = flag.String("tls-authority-cert", "", "Optional: Create TLS certificate/keys based on this CA authority. Useful for integration testing.")
 	authorityKeyFile  = flag.String("tls-authority-key", "", "Optional: Create TLS certificate/keys based on this CA authority. Useful for integration testing.")
 	serverName        = flag.String("server", "", "Optional: Create TLS certificate/keys with this domain name set. Useful for integration testing.")
+	keySize           = flag.Int("keysize", 4096, "Optional: Create TLS RSA private key with the given key size")
 )
 
 func main() {
@@ -58,12 +59,12 @@ func main() {
 			log.Fatal("Zero or both of --tls-key and --tls-cert must be supplied")
 		}
 		if *authorityCertFile == "" && *authorityKeyFile == "" {
-			if err := test.NewTLSKey(*tlsKeyFile, *tlsCertFile); err != nil {
+			if err := test.NewTLSKey(*tlsKeyFile, *tlsCertFile, *keySize); err != nil {
 				panic(err)
 			}
 		} else {
 			// generate the TLS cert/key based on the authority given.
-			if err := test.NewTLSKeyWithAuthority(*serverName, *tlsKeyFile, *tlsCertFile, *authorityKeyFile, *authorityCertFile); err != nil {
+			if err := test.NewTLSKeyWithAuthority(*serverName, *tlsKeyFile, *tlsCertFile, *authorityKeyFile, *authorityCertFile, *keySize); err != nil {
 				panic(err)
 			}
 		}
diff --git a/test/http.go b/test/http.go
index 37b3648f8..84207762d 100644
--- a/test/http.go
+++ b/test/http.go
@@ -68,7 +68,7 @@ func ListenAndServe(t *testing.T, router http.Handler, withTLS bool) (apiURL str
 		if withTLS {
 			certFile := filepath.Join(t.TempDir(), "dendrite.cert")
 			keyFile := filepath.Join(t.TempDir(), "dendrite.key")
-			err = NewTLSKey(keyFile, certFile)
+			err = NewTLSKey(keyFile, certFile, 64)
 			if err != nil {
 				t.Errorf("failed to make TLS key: %s", err)
 				return
diff --git a/test/keys.go b/test/keys.go
index 327c6ed7b..fb156ef27 100644
--- a/test/keys.go
+++ b/test/keys.go
@@ -69,8 +69,8 @@ func NewMatrixKey(matrixKeyPath string) (err error) {
 
 const certificateDuration = time.Hour * 24 * 365 * 10
 
-func generateTLSTemplate(dnsNames []string) (*rsa.PrivateKey, *x509.Certificate, error) {
-	priv, err := rsa.GenerateKey(rand.Reader, 4096)
+func generateTLSTemplate(dnsNames []string, bitSize int) (*rsa.PrivateKey, *x509.Certificate, error) {
+	priv, err := rsa.GenerateKey(rand.Reader, bitSize)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -118,8 +118,8 @@ func writePrivateKey(tlsKeyPath string, priv *rsa.PrivateKey) error {
 }
 
 // NewTLSKey generates a new RSA TLS key and certificate and writes it to a file.
-func NewTLSKey(tlsKeyPath, tlsCertPath string) error {
-	priv, template, err := generateTLSTemplate(nil)
+func NewTLSKey(tlsKeyPath, tlsCertPath string, keySize int) error {
+	priv, template, err := generateTLSTemplate(nil, keySize)
 	if err != nil {
 		return err
 	}
@@ -136,8 +136,8 @@ func NewTLSKey(tlsKeyPath, tlsCertPath string) error {
 	return writePrivateKey(tlsKeyPath, priv)
 }
 
-func NewTLSKeyWithAuthority(serverName, tlsKeyPath, tlsCertPath, authorityKeyPath, authorityCertPath string) error {
-	priv, template, err := generateTLSTemplate([]string{serverName})
+func NewTLSKeyWithAuthority(serverName, tlsKeyPath, tlsCertPath, authorityKeyPath, authorityCertPath string, keySize int) error {
+	priv, template, err := generateTLSTemplate([]string{serverName}, keySize)
 	if err != nil {
 		return err
 	}