From 80e6ca5f05758949f41883e9f201464986e38422 Mon Sep 17 00:00:00 2001
From: Daniel Aloni <daniel@globekeeper.com>
Date: Mon, 18 Sep 2023 14:52:40 +0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=B8=20Update=20`LatestKeysUploadTs`=20?=
 =?UTF-8?q?on=20Cross=20Signing=20Keys=20Uploads?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 clientapi/routing/key_crosssigning.go | 49 +++++++++++++++++++++++++++
 userapi/api/api.go                    |  6 ++++
 2 files changed, 55 insertions(+)

diff --git a/clientapi/routing/key_crosssigning.go b/clientapi/routing/key_crosssigning.go
index a6c7958c4..200db1a41 100644
--- a/clientapi/routing/key_crosssigning.go
+++ b/clientapi/routing/key_crosssigning.go
@@ -15,7 +15,10 @@
 package routing
 
 import (
+	"encoding/json"
+	"fmt"
 	"net/http"
+	"time"
 
 	"github.com/matrix-org/dendrite/clientapi/auth"
 	"github.com/matrix-org/dendrite/clientapi/auth/authtypes"
@@ -99,6 +102,52 @@ func UploadCrossSigningDeviceKeys(
 		}
 	}
 
+	// Following additional logic is implemented to follow the [Notion PRD](https://globekeeper.notion.site/Account-Data-State-Event-c64c8df8025a494d86d3137d4e080ece)
+	if device.UserID != "" {
+		prevAccountDataReq := api.QueryAccountDataRequest{
+			UserID:   device.UserID,
+			DataType: "account_data",
+			RoomID:   "",
+		}
+		accountDataRes := api.QueryAccountDataResponse{}
+		if err := accountAPI.QueryAccountData(req.Context(), &prevAccountDataReq, &accountDataRes); err != nil {
+			util.GetLogger(req.Context()).WithError(err).Error("userAPI.QueryAccountData failed")
+			return util.ErrorResponse(fmt.Errorf("userAPI.QueryAccountData: %w", err))
+		}
+		var accoundData api.AccountData
+		if len(accountDataRes.GlobalAccountData) != 0 {
+			err := json.Unmarshal(accountDataRes.GlobalAccountData["account_data"], &accoundData)
+			if err != nil {
+				return util.JSONResponse{
+					Code: http.StatusInternalServerError,
+					JSON: spec.InternalServerError{Err: err.Error()},
+				}
+			}
+		}
+		accoundData.LatestKeysUploadTs = time.Now().UnixMilli()
+		newAccountData, err := json.Marshal(accoundData)
+		if err != nil {
+			return util.JSONResponse{
+				Code: http.StatusInternalServerError,
+				JSON: spec.InternalServerError{Err: err.Error()},
+			}
+		}
+
+		dataReq := api.InputAccountDataRequest{
+			UserID:      device.UserID,
+			DataType:    "account_data",
+			RoomID:      "",
+			AccountData: json.RawMessage(newAccountData),
+		}
+		dataRes := api.InputAccountDataResponse{}
+		if err := accountAPI.InputAccountData(req.Context(), &dataReq, &dataRes); err != nil {
+			util.GetLogger(req.Context()).WithError(err).Error("userAPI.InputAccountData on LatestKeysUploadTs update failed")
+			return util.ErrorResponse(err)
+		}
+		logger := util.GetLogger(req.Context()).WithField("user_id", device.UserID)
+		logger.Info("updated latestKeysUploadTs field in account data")
+	}
+
 	return util.JSONResponse{
 		Code: http.StatusOK,
 		JSON: struct{}{},
diff --git a/userapi/api/api.go b/userapi/api/api.go
index a072903a5..65ffca63b 100644
--- a/userapi/api/api.go
+++ b/userapi/api/api.go
@@ -296,6 +296,12 @@ type QueryAccountDataResponse struct {
 	RoomAccountData   map[string]map[string]json.RawMessage // room -> type -> data
 }
 
+// Custom Connnect AccountData information
+type AccountData struct {
+	IsProfileFilled    bool  `json:"isProfileFilled"`
+	LatestKeysUploadTs int64 `json:"latestKeysUploadTs"`
+}
+
 // QueryDevicesRequest is the request for QueryDevices
 type QueryDevicesRequest struct {
 	UserID string