Improve key import/export

crypto/keyexport.go

@@ -7,6 +7,7 @@
 package crypto
 
 import (
+	"bytes"
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/hmac"
@@ -16,12 +17,13 @@ import (
 	"encoding/base64"
 	"encoding/binary"
 	"encoding/json"
+	"fmt"
 	"math"
-	"strings"
 
 	"github.com/pkg/errors"
 	"golang.org/x/crypto/pbkdf2"
 
+	"maunium.net/go/mautrix/crypto/olm"
 	"maunium.net/go/mautrix/id"
 )
 
@@ -39,8 +41,20 @@ type ExportedSession struct {
 	SessionKey        string            `json:"session_key"`
 }
 
+// The default number of pbkdf2 rounds to use when exporting keys
 const defaultPassphraseRounds = 100000
 
+const exportPrefix = "-----BEGIN MEGOLM SESSION DATA-----\n"
+const exportSuffix = "-----END MEGOLM SESSION DATA-----\n"
+// Only version 0x01 is currently specified in the spec
+const exportVersion1 = 0x01
+// The standard for wrapping base64 is 76 bytes
+const exportLineLengthLimit = 76
+// Byte count for version + salt + iv + number of rounds
+const exportHeaderLength = 1 + 16 + 16 + 4
+// SHA-256 hash length
+const exportHashLength = 32
+
 func computeKey(passphrase string, salt []byte, rounds int) (encryptionKey, hashKey []byte) {
 	key := pbkdf2.Key([]byte(passphrase), salt, rounds, 64, sha512.New)
 	encryptionKey = key[:32]
@@ -48,27 +62,27 @@ func computeKey(passphrase string, salt []byte, rounds int) (encryptionKey, hash
 	return
 }
 
-func makeExportIV() ([]byte, error) {
+func makeExportIV() []byte {
 	iv := make([]byte, 16)
 	_, err := rand.Read(iv)
 	if err != nil {
-		return nil, err
+		panic(olm.NotEnoughGoRandom)
 	}
 	// Set bit 63 to zero
 	iv[7] &= 0b11111110
-	return iv, nil
+	return iv
 }
 
-func makeExportKeys(passphrase string) (encryptionKey, hashKey, salt, iv []byte, err error) {
+func makeExportKeys(passphrase string) (encryptionKey, hashKey, salt, iv []byte) {
 	salt = make([]byte, 16)
-	_, err = rand.Read(salt)
+	_, err := rand.Read(salt)
 	if err != nil {
-		return
+		panic(olm.NotEnoughGoRandom)
 	}
 
 	encryptionKey, hashKey = computeKey(passphrase, salt, defaultPassphraseRounds)
 
-	iv, err = makeExportIV()
+	iv = makeExportIV()
 	return
 }
 
@@ -100,12 +114,6 @@ func exportSessionsJSON(sessions []*InboundGroupSession) ([]byte, error) {
 	return json.Marshal(exportedSessions)
 }
 
-const exportPrefix = "-----BEGIN MEGOLM SESSION DATA-----"
-const exportSuffix = "-----END MEGOLM SESSION DATA-----"
-const exportLineLengthLimit = 76
-const exportHeaderLength = 1 + 16 + 16 + 4
-const exportHashLength = 32
-
 func min(a, b int) int {
 	if a > b {
 		return b
@@ -113,32 +121,38 @@ func min(a, b int) int {
 	return a
 }
 
-func formatKeyExportData(data []byte) string {
-	dataStr := base64.StdEncoding.EncodeToString(data)
-	// Base64 lines + prefix + suffix + empty line at end
-	lines := make([]string, int(math.Ceil(float64(len(dataStr)) / exportLineLengthLimit)) + 3)
-	lines[0] = exportPrefix
-	line := 1
-	for ptr := 0; ptr < len(dataStr); ptr += exportLineLengthLimit {
-		lines[line] = dataStr[ptr:min(ptr+exportLineLengthLimit, len(dataStr))]
-		line++
+func formatKeyExportData(data []byte) []byte {
+	base64Data := make([]byte, base64.StdEncoding.EncodedLen(len(data)))
+	base64.StdEncoding.Encode(base64Data, data)
+
+	// Prefix + data and newline for each 76 characters of data + suffix
+	outputLength := len(exportPrefix) +
+		len(base64Data) + int(math.Ceil(float64(len(base64Data))/exportLineLengthLimit)) +
+		len(exportSuffix)
+
+	var buf bytes.Buffer
+	buf.Grow(outputLength)
+	buf.WriteString(exportPrefix)
+	for ptr := 0; ptr < len(base64Data); ptr += exportLineLengthLimit {
+		buf.Write(base64Data[ptr:min(ptr+exportLineLengthLimit, len(base64Data))])
+		buf.WriteRune('\n')
 	}
-	lines[len(lines)-2] = exportSuffix
-	return strings.Join(lines, "\n")
+	buf.WriteString(exportSuffix)
+	if buf.Len() != buf.Cap() || buf.Len() != outputLength {
+		panic(fmt.Errorf("unexpected length %d / %d / %d", buf.Len(), buf.Cap(), outputLength))
+	}
+	return buf.Bytes()
 }
 
 // ExportKeys exports the given Megolm sessions with the format specified in the Matrix spec.
 // See https://matrix.org/docs/spec/client_server/r0.6.1#key-exports
-func ExportKeys(passphrase string, sessions []*InboundGroupSession) (string, error) {
+func ExportKeys(passphrase string, sessions []*InboundGroupSession) ([]byte, error) {
 	// Make all the keys necessary for exporting
-	encryptionKey, hashKey, salt, iv, err := makeExportKeys(passphrase)
-	if err != nil {
-		return "", err
-	}
+	encryptionKey, hashKey, salt, iv := makeExportKeys(passphrase)
 	// Export all the given sessions and put them in JSON
 	unencryptedData, err := exportSessionsJSON(sessions)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 
 	// The export data consists of:
@@ -153,16 +167,13 @@ func ExportKeys(passphrase string, sessions []*InboundGroupSession) (string, err
 	dataWithoutHashLength := len(exportData) - exportHashLength
 
 	// Create the header for the export data
-	exportData[0] = 0x01
+	exportData[0] = exportVersion1
 	copy(exportData[1:17], salt)
 	copy(exportData[17:33], iv)
 	binary.BigEndian.PutUint32(exportData[33:37], defaultPassphraseRounds)
 
 	// Encrypt data with AES-256-CTR
-	block, err := aes.NewCipher(encryptionKey)
-	if err != nil {
-		return "", err
-	}
+	block, _ := aes.NewCipher(encryptionKey)
 	cipher.NewCTR(block, iv).XORKeyStream(exportData[exportHeaderLength:dataWithoutHashLength], unencryptedData)
 
 	// Hash all the data with HMAC-SHA256 and put it at the end
@@ -172,4 +183,4 @@ func ExportKeys(passphrase string, sessions []*InboundGroupSession) (string, err
 
 	// Format the export (prefix, base64'd exportData, suffix) and return
 	return formatKeyExportData(exportData), nil
-}
+}