From 0cf8e9166cdcdb1bbf76bd09d2afa2f49b06b0ad Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Wed, 23 Sep 2020 10:20:35 +0100
Subject: [PATCH] Enforce valid key IDs

---
 internal/config/config.go | 6 ++++++
 internal/test/config.go   | 7 ++++++-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/internal/config/config.go b/internal/config/config.go
index d7470f873..d75500db5 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -36,6 +36,9 @@ import (
 	jaegermetrics "github.com/uber/jaeger-lib/metrics"
 )
 
+// keyIDRegexp defines allowable characters in Key IDs.
+var keyIDRegexp = regexp.MustCompile("^ed25519:[a-zA-Z0-9_]+$")
+
 // Version is the current version of the config format.
 // This will change whenever we make breaking changes to the config format.
 const Version = 1
@@ -459,6 +462,9 @@ func readKeyPEM(path string, data []byte) (gomatrixserverlib.KeyID, ed25519.Priv
 			if !strings.HasPrefix(keyID, "ed25519:") {
 				return "", nil, fmt.Errorf("key ID %q doesn't start with \"ed25519:\" in %q", keyID, path)
 			}
+			if !keyIDRegexp.MatchString(keyID) {
+				return "", nil, fmt.Errorf("key ID %q in %q contains illegal characters (use a-z, A-Z, 0-9 and _ only)", keyID, path)
+			}
 			_, privKey, err := ed25519.GenerateKey(bytes.NewReader(keyBlock.Bytes))
 			if err != nil {
 				return "", nil, err
diff --git a/internal/test/config.go b/internal/test/config.go
index 72cd0e6e4..8080988f3 100644
--- a/internal/test/config.go
+++ b/internal/test/config.go
@@ -25,6 +25,7 @@ import (
 	"math/big"
 	"os"
 	"path/filepath"
+	"strings"
 	"time"
 
 	"github.com/matrix-org/dendrite/internal/config"
@@ -146,10 +147,14 @@ func NewMatrixKey(matrixKeyPath string) (err error) {
 		err = keyOut.Close()
 	})()
 
+	keyID := base64.RawURLEncoding.EncodeToString(data[:])
+	keyID = strings.ReplaceAll(keyID, "-", "")
+	keyID = strings.ReplaceAll(keyID, "_", "")
+
 	err = pem.Encode(keyOut, &pem.Block{
 		Type: "MATRIX PRIVATE KEY",
 		Headers: map[string]string{
-			"Key-ID": "ed25519:" + base64.RawStdEncoding.EncodeToString(data[:3]),
+			"Key-ID": fmt.Sprintf("ed25519:%s", keyID[:6]),
 		},
 		Bytes: data[3:],
 	})