crypto: fix key exports

crypto/keyexport.go

@@ -1,4 +1,4 @@
-// Copyright (c) 2020 Tulir Asokan
+// Copyright (c) 2025 Tulir Asokan
 //
 // This Source Code Form is subject to the terms of the Mozilla Public
 // License, v. 2.0. If a copy of the MPL was not distributed with this
@@ -16,15 +16,21 @@ import (
 	"encoding/base64"
 	"encoding/binary"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"math"
 
+	"go.mau.fi/util/dbutil"
+	"go.mau.fi/util/exbytes"
+	"go.mau.fi/util/exerrors"
 	"go.mau.fi/util/random"
 	"golang.org/x/crypto/pbkdf2"
 
 	"maunium.net/go/mautrix/id"
 )
 
+var ErrNoSessionsForExport = errors.New("no sessions provided for export")
+
 type SenderClaimedKeys struct {
 	Ed25519 id.Ed25519 `json:"ed25519"`
 }
@@ -78,22 +84,14 @@ func makeExportKeys(passphrase string) (encryptionKey, hashKey, salt, iv []byte)
 	return
 }
 
-func exportSessions(sessions []*InboundGroupSession) ([]ExportedSession, error) {
-	export := make([]ExportedSession, len(sessions))
+func exportSessions(sessions []*InboundGroupSession) ([]*ExportedSession, error) {
+	export := make([]*ExportedSession, len(sessions))
+	var err error
 	for i, session := range sessions {
-		key, err := session.Internal.Export(session.Internal.FirstKnownIndex())
+		export[i], err = session.export()
 		if err != nil {
 			return nil, fmt.Errorf("failed to export session: %w", err)
 		}
-		export[i] = ExportedSession{
-			Algorithm:         id.AlgorithmMegolmV1,
-			ForwardingChains:  session.ForwardingChains,
-			RoomID:            session.RoomID,
-			SenderKey:         session.SenderKey,
-			SenderClaimedKeys: SenderClaimedKeys{},
-			SessionID:         session.ID(),
-			SessionKey:        string(key),
-		}
 	}
 	return export, nil
 }
@@ -107,38 +105,73 @@ func exportSessionsJSON(sessions []*InboundGroupSession) ([]byte, error) {
 }
 
 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
+	encodedLen := base64.StdEncoding.EncodedLen(len(data))
 	outputLength := len(exportPrefix) +
-		len(base64Data) + int(math.Ceil(float64(len(base64Data))/exportLineLengthLimit)) +
+		encodedLen + int(math.Ceil(float64(encodedLen)/exportLineLengthLimit)) +
 		len(exportSuffix)
+	output := make([]byte, 0, outputLength)
+	outputWriter := (*exbytes.Writer)(&output)
+	base64Writer := base64.NewEncoder(base64.StdEncoding, outputWriter)
+	lineByteCount := base64.StdEncoding.DecodedLen(exportLineLengthLimit)
+	exerrors.Must(outputWriter.WriteString(exportPrefix))
+	for i := 0; i < len(data); i += lineByteCount {
+		exerrors.Must(base64Writer.Write(data[i:min(i+lineByteCount, len(data))]))
+		if i+lineByteCount >= len(data) {
+			exerrors.PanicIfNotNil(base64Writer.Close())
+		}
+		exerrors.PanicIfNotNil(outputWriter.WriteByte('\n'))
+	}
+	exerrors.Must(outputWriter.WriteString(exportSuffix))
+	if len(output) != outputLength {
+		panic(fmt.Errorf("unexpected length %d / %d", len(output), outputLength))
+	}
+	return output
+}
 
-	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')
+func ExportKeysIter(passphrase string, sessions dbutil.RowIter[*InboundGroupSession]) ([]byte, error) {
+	buf := bytes.NewBuffer(make([]byte, 0, 50*1024))
+	enc := json.NewEncoder(buf)
+	buf.WriteByte('[')
+	err := sessions.Iter(func(session *InboundGroupSession) (bool, error) {
+		exported, err := session.export()
+		if err != nil {
+			return false, err
+		}
+		err = enc.Encode(exported)
+		if err != nil {
+			return false, err
+		}
+		buf.WriteByte(',')
+		return true, nil
+	})
+	if err != nil {
+		return nil, err
 	}
-	buf.WriteString(exportSuffix)
-	if buf.Len() != buf.Cap() || buf.Len() != outputLength {
-		panic(fmt.Errorf("unexpected length %d / %d / %d", buf.Len(), buf.Cap(), outputLength))
+	output := buf.Bytes()
+	if len(output) == 1 {
+		return nil, ErrNoSessionsForExport
 	}
-	return buf.Bytes()
+	output[len(output)-1] = ']' // Replace the last comma with a closing bracket
+	return EncryptKeyExport(passphrase, output)
 }
 
 // ExportKeys exports the given Megolm sessions with the format specified in the Matrix spec.
 // See https://spec.matrix.org/v1.2/client-server-api/#key-exports
 func ExportKeys(passphrase string, sessions []*InboundGroupSession) ([]byte, error) {
-	// Make all the keys necessary for exporting
-	encryptionKey, hashKey, salt, iv := makeExportKeys(passphrase)
+	if len(sessions) == 0 {
+		return nil, ErrNoSessionsForExport
+	}
 	// Export all the given sessions and put them in JSON
 	unencryptedData, err := exportSessionsJSON(sessions)
 	if err != nil {
 		return nil, err
 	}
+	return EncryptKeyExport(passphrase, unencryptedData)
+}
+
+func EncryptKeyExport(passphrase string, unencryptedData json.RawMessage) ([]byte, error) {
+	// Make all the keys necessary for exporting
+	encryptionKey, hashKey, salt, iv := makeExportKeys(passphrase)
 
 	// The export data consists of:
 	// 1 byte of export format version

crypto/keyexport_test.go

@@ -0,0 +1,35 @@
+// Copyright (c) 2024 Tulir Asokan
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+package crypto_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"go.mau.fi/util/exerrors"
+	"go.mau.fi/util/exfmt"
+
+	"maunium.net/go/mautrix/crypto"
+	"maunium.net/go/mautrix/crypto/olm"
+)
+
+func TestExportKeys(t *testing.T) {
+	acc := crypto.NewOlmAccount()
+	sess := exerrors.Must(crypto.NewInboundGroupSession(
+		acc.IdentityKey(),
+		acc.SigningKey(),
+		"!room:example.com",
+		exerrors.Must(olm.NewOutboundGroupSession()).Key(),
+		7*exfmt.Day,
+		100,
+		false,
+	))
+	data, err := crypto.ExportKeys("meow", []*crypto.InboundGroupSession{sess})
+	assert.NoError(t, err)
+	assert.Len(t, data, 840)
+}

crypto/keyimport.go

@@ -36,6 +36,10 @@ var (
 var exportPrefixBytes, exportSuffixBytes = []byte(exportPrefix), []byte(exportSuffix)
 
 func decodeKeyExport(data []byte) ([]byte, error) {
+	// Fix some types of corruption in the key export file before checking anything
+	if bytes.IndexByte(data, '\r') != -1 {
+		data = bytes.ReplaceAll(data, []byte{'\r', '\n'}, []byte{'\n'})
+	}
 	// If the valid prefix and suffix aren't there, it's probably not a Matrix key export
 	if !bytes.HasPrefix(data, exportPrefixBytes) {
 		return nil, ErrMissingExportPrefix

crypto/sessions.go

@@ -8,6 +8,7 @@ package crypto
 
 import (
 	"errors"
+	"fmt"
 	"time"
 
 	"maunium.net/go/mautrix/crypto/olm"
@@ -152,6 +153,22 @@ func (igs *InboundGroupSession) RatchetTo(index uint32) error {
 	return nil
 }
 
+func (igs *InboundGroupSession) export() (*ExportedSession, error) {
+	key, err := igs.Internal.Export(igs.Internal.FirstKnownIndex())
+	if err != nil {
+		return nil, fmt.Errorf("failed to export session: %w", err)
+	}
+	return &ExportedSession{
+		Algorithm:         id.AlgorithmMegolmV1,
+		ForwardingChains:  igs.ForwardingChains,
+		RoomID:            igs.RoomID,
+		SenderKey:         igs.SenderKey,
+		SenderClaimedKeys: SenderClaimedKeys{},
+		SessionID:         igs.ID(),
+		SessionKey:        string(key),
+	}, nil
+}
+
 type OGSState int
 
 const (

go.mod

@@ -18,7 +18,7 @@ require (
 	github.com/tidwall/gjson v1.18.0
 	github.com/tidwall/sjson v1.2.5
 	github.com/yuin/goldmark v1.7.8
-	go.mau.fi/util v0.8.5-0.20250129121406-18c356e558b8
+	go.mau.fi/util v0.8.5-0.20250203220331-1c0d19ea6003
 	go.mau.fi/zeroconfig v0.1.3
 	golang.org/x/crypto v0.32.0
 	golang.org/x/exp v0.0.0-20250128182459-e0ece0dbea4c

go.sum

@@ -54,8 +54,8 @@ github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY=
 github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28=
 github.com/yuin/goldmark v1.7.8 h1:iERMLn0/QJeHFhxSt3p6PeN9mGnvIKSpG9YYorDMnic=
 github.com/yuin/goldmark v1.7.8/go.mod h1:uzxRWxtg69N339t3louHJ7+O03ezfj6PlliRlaOzY1E=
-go.mau.fi/util v0.8.5-0.20250129121406-18c356e558b8 h1:O1cRlXPahwbu1ckIf8XgUP3gHMJlSqJxaVTqwRlVK4s=
-go.mau.fi/util v0.8.5-0.20250129121406-18c356e558b8/go.mod h1:MOfGTs1CBuK6ERTcSL4lb5YU7/ujz09eOPVEDckuazY=
+go.mau.fi/util v0.8.5-0.20250203220331-1c0d19ea6003 h1:ye5l+QpYW5CpGVMedb3EHlmflGMQsMtw8mC4K/U8hIw=
+go.mau.fi/util v0.8.5-0.20250203220331-1c0d19ea6003/go.mod h1:MOfGTs1CBuK6ERTcSL4lb5YU7/ujz09eOPVEDckuazY=
 go.mau.fi/zeroconfig v0.1.3 h1:As9wYDKmktjmNZW5i1vn8zvJlmGKHeVxHVIBMXsm4kM=
 go.mau.fi/zeroconfig v0.1.3/go.mod h1:NcSJkf180JT+1IId76PcMuLTNa1CzsFFZ0nBygIQM70=
 golang.org/x/crypto v0.32.0 h1:euUpcYgM8WcP71gNpTqQCn6rC2t6ULUPiOzfWaXVVfc=