From ec5d1d681d1362f5746c5cb45e93829d6a68aa4d Mon Sep 17 00:00:00 2001
From: Till <2353100+S7evinK@users.noreply.github.com>
Date: Thu, 6 Oct 2022 12:30:24 +0200
Subject: [PATCH] Always return `one_time_key_counts` on `/keys/upload` (#2769)

The OTK count is
[required](https://spec.matrix.org/v1.4/client-server-api/#post_matrixclientv3keysupload)
in responses to `/keys/upload`, so return those.
---
 clientapi/routing/keys.go      | 4 ++--
 keyserver/internal/internal.go | 5 +++++
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/clientapi/routing/keys.go b/clientapi/routing/keys.go
index b7a76b47..5c368138 100644
--- a/clientapi/routing/keys.go
+++ b/clientapi/routing/keys.go
@@ -19,11 +19,12 @@ import (
 	"net/http"
 	"time"
 
+	"github.com/matrix-org/util"
+
 	"github.com/matrix-org/dendrite/clientapi/httputil"
 	"github.com/matrix-org/dendrite/clientapi/jsonerror"
 	"github.com/matrix-org/dendrite/keyserver/api"
 	userapi "github.com/matrix-org/dendrite/userapi/api"
-	"github.com/matrix-org/util"
 )
 
 type uploadKeysRequest struct {
@@ -77,7 +78,6 @@ func UploadKeys(req *http.Request, keyAPI api.ClientKeyAPI, device *userapi.Devi
 		}
 	}
 	keyCount := make(map[string]int)
-	// we only return key counts when the client uploads OTKs
 	if len(uploadRes.OneTimeKeyCounts) > 0 {
 		keyCount = uploadRes.OneTimeKeyCounts[0].KeyCount
 	}
diff --git a/keyserver/internal/internal.go b/keyserver/internal/internal.go
index 017c29e8..a0280dff 100644
--- a/keyserver/internal/internal.go
+++ b/keyserver/internal/internal.go
@@ -70,6 +70,11 @@ func (a *KeyInternalAPI) PerformUploadKeys(ctx context.Context, req *api.Perform
 	if len(req.OneTimeKeys) > 0 {
 		a.uploadOneTimeKeys(ctx, req, res)
 	}
+	otks, err := a.DB.OneTimeKeysCount(ctx, req.UserID, req.DeviceID)
+	if err != nil {
+		return err
+	}
+	res.OneTimeKeyCounts = []api.OneTimeKeysCount{*otks}
 	return nil
 }