Add OTK re-upload test

TEST_HITLIST.md

@@ -26,7 +26,7 @@ Key Verification: (Short Authentication String)
 - [ ] Verification can be cancelled over federation.
 
 Network connectivity:
-- [ ] If a client cannot upload OTKs, it retries.
+- [x] If a client cannot upload OTKs, it retries.
 - [ ] If a client cannot claim local OTKs, it retries.
 - [ ] If a client cannot claim remote OTKs, it retries.
 - [x] If a server cannot send device list updates over federation, it retries. https://github.com/matrix-org/complement/pull/695

tests/main_test.go

@@ -55,6 +55,15 @@ func ClientTypeMatrix(t *testing.T, subTest func(tt *testing.T, a, b api.ClientT
 	}
 }
 
+func ForEachClientType(t *testing.T, subTest func(tt *testing.T, a api.ClientType)) {
+	for _, tc := range []api.ClientType{{Lang: api.ClientTypeRust, HS: "hs1"}, {Lang: api.ClientTypeJS, HS: "hs1"}} {
+		tc := tc
+		t.Run(string(tc.Lang), func(t *testing.T) {
+			subTest(t, tc)
+		})
+	}
+}
+
 func MustCreateClient(t *testing.T, clientType api.ClientType, cfg api.ClientCreationOpts, ssURL string, opts ...func(api.Client, api.ClientCreationOpts)) api.Client {
 	var c api.Client
 	switch clientType.Lang {
@@ -87,22 +96,28 @@ type TestContext struct {
 	Bob        *client.CSAPI
 }
 
-func CreateTestContext(t *testing.T, clientTypeA, clientTypeB api.ClientType) *TestContext {
+func CreateTestContext(t *testing.T, clientType ...api.ClientType) *TestContext {
 	deployment := Deploy(t)
-	// pre-register alice and bob
-	csapiAlice := deployment.Register(t, clientTypeA.HS, helpers.RegistrationOpts{
-		LocalpartSuffix: "alice",
-		Password:        "complement-crypto-password",
-	})
-	csapiBob := deployment.Register(t, clientTypeB.HS, helpers.RegistrationOpts{
-		LocalpartSuffix: "bob",
-		Password:        "complement-crypto-password",
-	})
-	return &TestContext{
+	tc := &TestContext{
 		Deployment: deployment,
-		Alice:      csapiAlice,
-		Bob:        csapiBob,
 	}
+	// pre-register alice and bob, if told
+	if len(clientType) > 0 {
+		tc.Alice = deployment.Register(t, clientType[0].HS, helpers.RegistrationOpts{
+			LocalpartSuffix: "alice",
+			Password:        "complement-crypto-password",
+		})
+	}
+	if len(clientType) > 1 {
+		tc.Bob = deployment.Register(t, clientType[1].HS, helpers.RegistrationOpts{
+			LocalpartSuffix: "bob",
+			Password:        "complement-crypto-password",
+		})
+	}
+	if len(clientType) > 2 {
+		t.Fatalf("CreateTestContext: too many clients: got %d", len(clientType))
+	}
+	return tc
 }
 
 func (c *TestContext) CreateNewEncryptedRoom(t *testing.T, creator *client.CSAPI, preset string, invite []string) (roomID string) {

tests/mitmproxy_addons/status_code.py

@@ -12,19 +12,25 @@ def __init__(self):
         print(MITM_DOMAIN_NAME)
         self.matchall = flowfilter.parse(".")
         self.filter: Optional[flowfilter.TFilter] = self.matchall
+        self._seen_count = 0
 
     def reset(self):
         self.config = {
             "return_status": 0,
             "block_request": False,
+            "count": 0,
             "filter": None,
         }
+        self._seen_count = 0
+
+    def exceeded_count(self):
+        return self.config["count"] > 0 and self._seen_count >= self.config["count"]
 
     def load(self, loader):
         loader.add_option(
             name="statuscode",
             typespec=dict,
-            default={"return_status": 0, "filter": None, "block_request": False},
+            default={"return_status": 0, "filter": None, "block_request": False, "count": 0},
             help="Change the response status code, with an optional filter",
         )
 
@@ -37,7 +43,7 @@ def configure(self, updates):
             return
         self.config = ctx.options.statuscode
         new_filter = self.config.get('filter', None)
-        print(f"statuscode will return HTTP {self.config['return_status']} filter={new_filter}")
+        print(f"statuscode will return HTTP {self.config['return_status']} filter={new_filter} count={self.config['count']}")
         if new_filter:
             self.filter = flowfilter.parse(new_filter)
         else:
@@ -49,15 +55,23 @@ def request(self, flow):
             return
         if self.config["return_status"] == 0:
             return # ignore responses if we aren't told a code
-        if self.config["block_request"] and flowfilter.match(self.filter, flow):
-            print(f'statuscode: blocking request and sending back {self.config["return_status"]}')
-            flow.response = Response.make(self.config["return_status"])
+        if self.config["block_request"] and flowfilter.match(self.filter, flow) and not self.exceeded_count():
+            self._seen_count += 1
+            print(f'statuscode: blocking request and sending back {self.config["return_status"]}: count {self._seen_count}/{self.config["count"]}')
+            flow.response = Response.make(self.config["return_status"], headers={
+                "MITM-Proxy": "yes",
+            })
 
     def response(self, flow):
         # always ignore the controller
         if flow.request.pretty_host == MITM_DOMAIN_NAME:
             return
         if self.config["return_status"] == 0:
             return # ignore responses if we aren't told a code
-        if flowfilter.match(self.filter, flow):
+        if flow.response.headers.get("MITM-Proxy", None) is not None:
+            print(f'not modifying mitm response!')
+            return # ignore responses generated by mitm proxy (i.e the one in `def request` above
+        if flowfilter.match(self.filter, flow) and not self.exceeded_count():
+            self._seen_count += 1
+            print(f'statuscode: blocking response and sending back {self.config["return_status"]}: count {self._seen_count}/{self.config["count"]}')
             flow.response = Response.make(self.config["return_status"])

tests/one_time_keys_test.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/matrix-org/complement-crypto/internal/api"
+	"github.com/matrix-org/complement/b"
 	"github.com/matrix-org/complement/client"
 	"github.com/matrix-org/complement/ct"
 	"github.com/matrix-org/complement/helpers"
@@ -182,3 +183,61 @@ func TestFallbackKeyIsUsedIfOneTimeKeysRunOut(t *testing.T) {
 
 	})
 }
+
+func TestFailedOneTimeKeyUploadRetries(t *testing.T) {
+	ForEachClientType(t, func(t *testing.T, clientType api.ClientType) {
+		tc := CreateTestContext(t, clientType, clientType)
+		// make a room so we can kick clients
+		roomID := tc.Alice.MustCreateRoom(t, map[string]interface{}{"preset": "public_chat"})
+		// block /keys/upload and make a client
+		tc.Deployment.WithMITMOptions(t, map[string]interface{}{
+			"statuscode": map[string]interface{}{
+				"return_status": http.StatusGatewayTimeout,
+				"block_request": true,
+				"count":         2, // block it twice
+				"filter":        "~u .*\\/keys\\/upload.*",
+			},
+		}, func() {
+			alice := LoginClientFromComplementClient(t, tc.Deployment, tc.Alice, clientType)
+			defer alice.Close(t)
+			aliceStopSyncing := alice.MustStartSyncing(t)
+			defer aliceStopSyncing()
+
+			// we should be able to claim a key eventually
+			for i := 0; i < 5; i++ {
+				time.Sleep(time.Duration(200 * time.Millisecond * time.Duration(i+1)))
+				res := tc.Bob.MustDo(t, "POST", []string{
+					"_matrix", "client", "v3", "keys", "claim",
+				}, client.WithJSONBody(t, map[string]any{
+					"one_time_keys": map[string]any{
+						tc.Alice.UserID: map[string]any{
+							tc.Alice.DeviceID: "signed_curve25519",
+						},
+					},
+				}))
+				jsonBody := must.ParseJSON(t, res.Body)
+				res.Body.Close()
+				err := match.JSONKeyPresent(
+					fmt.Sprintf("one_time_keys.%s.%s.signed_curve25519*", tc.Alice.UserID, tc.Alice.DeviceID),
+				)(jsonBody)
+				if err == nil {
+					break
+				}
+				t.Logf("failed to claim otk: /keys/claim => %v", jsonBody.Raw)
+				if i == 4 {
+					t.Errorf("failed to claim OTK for user, did /keys/upload retry?")
+				}
+
+				// try kicking the client by sending some data down /sync
+				// Specifically, JS SDK needs this. Rust has its own backoff independent to /sync
+				tc.Alice.SendEventSynced(t, roomID, b.Event{
+					Type: "m.room.message",
+					Content: map[string]interface{}{
+						"msgtype": "m.text",
+						"body":    "this is a kick to try to get clients to retry /keys/upload",
+					},
+				})
+			}
+		})
+	})
+}