Merge pull request #106 from matrix-org/kegan/callback-req2

Remove status_code addon

internal/deploy/mitm.go

@@ -8,6 +8,7 @@ import (
 	"net/url"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -111,39 +112,48 @@ func (c *MITMConfiguration) Execute(inner func()) {
 	//       $addon_values...
 	//     }
 	//   }
-	// We have 2 add-ons: "callback" and "status_code". The former just sniffs
-	// requests/responses. The latter modifies the request/response in some way.
 	//
 	// The API shape of the add-ons are located inside the python files in tests/mitmproxy_addons
-	body := map[string]any{}
 	if len(c.pathCfgs) > 1 {
 		c.t.Fatalf(">1 path config currently unsupported") // TODO
 	}
 	c.mu.Lock()
+	callbackAddon := map[string]any{}
 	for _, pathConfig := range c.pathCfgs {
+		if pathConfig.filter() != "" {
+			callbackAddon["filter"] = pathConfig.filter()
+		}
+		cbServer, err := NewCallbackServer(c.t, c.client.hostnameRunningComplement)
+		must.NotError(c.t, "failed to start callback server", err)
+		defer cbServer.Close()
+
 		if pathConfig.listener != nil {
-			cbServer, err := NewCallbackServer(c.t, c.client.hostnameRunningComplement)
-			must.NotError(c.t, "failed to start callback server", err)
-			callbackURL := cbServer.SetOnResponseCallback(c.t, pathConfig.listener)
-			defer cbServer.Close()
-
-			body["callback"] = map[string]any{
-				"callback_response_url": callbackURL,
-				"filter":                pathConfig.filter(),
-			}
+			responseCallbackURL := cbServer.SetOnResponseCallback(c.t, pathConfig.listener)
+			callbackAddon["callback_response_url"] = responseCallbackURL
 		}
-		if pathConfig.blockRequest != nil {
-			body["statuscode"] = map[string]any{
-				"return_status": pathConfig.blockStatusCode,
-				"block_request": *pathConfig.blockRequest,
-				"count":         pathConfig.blockCount,
-				"filter":        pathConfig.filter(),
-			}
+		if pathConfig.blockRequest != nil && *pathConfig.blockRequest {
+			// reimplement statuscode plugin logic in Go
+			// TODO: refactor this
+			var count atomic.Uint32
+			requestCallbackURL := cbServer.SetOnRequestCallback(c.t, func(cd CallbackData) *CallbackResponse {
+				newCount := count.Add(1)
+				if pathConfig.blockCount > 0 && newCount > uint32(pathConfig.blockCount) {
+					return nil // don't block
+				}
+				// block this request by sending back a fake response
+				return &CallbackResponse{
+					RespondStatusCode: pathConfig.blockStatusCode,
+					RespondBody:       json.RawMessage(`{"error":"complement-crypto says no"}`),
+				}
+			})
+			callbackAddon["callback_request_url"] = requestCallbackURL
 		}
 	}
 	c.mu.Unlock()
 
-	lockID := c.client.lockOptions(c.t, body)
+	lockID := c.client.lockOptions(c.t, map[string]any{
+		"callback": callbackAddon,
+	})
 	defer c.client.unlockOptions(c.t, lockID)
 	inner()
 

tests/mitmproxy_addons/README.md

@@ -8,11 +8,8 @@ How this works:
  - It is also told to run a normal proxy, to which a Flask HTTP server is attached.
  - The Flask HTTP server can be used to control mitmproxy at test runtime. This is done via the Controller HTTP API.
 
-
 ### Controller HTTP API
 
-**This is highly experimental and will change without warning.**
-
 `mitmproxy` is run once for all tests. To avoid test pollution, the controller is "locked" for the duration
 of a test and must be "unlocked" afterwards. When acquiring the lock, options can be set on `mitmproxy`.
 
@@ -21,17 +18,100 @@ POST /options/lock
  {
    "options": {
      "body_size_limit": "3m",
+     "callback": {
+       "callback_response_url": "http://host.docker.internal:445566"
+     }
    }
  }
  HTTP/1.1 200 OK
  {
    "reset_id": "some_opaque_string"
  }
 ```
+Any [option](https://docs.mitmproxy.org/stable/concepts-options/) can be specified in the
+`options` object, not just Complement specific addons.
 
 ```
 POST /options/unlock
 {
    "reset_id": "some_opaque_string"
 }
-```
+```
+
+Tests will lock/unlock whenever they need to interact with mitmproxy. Attempting to lock an already locked controller will return an HTTP 400 error. Attempting to unlock an already unlocked controller will return an HTTP 400 error.
+
+### Callback addon
+
+A [mitmproxy addon](https://docs.mitmproxy.org/stable/addons-examples/) bolts on custom
+functionality to mitmproxy. This typically involves using the
+[Event Hooks API](https://docs.mitmproxy.org/stable/api/events.html) to listen for
+[HTTP flows](https://docs.mitmproxy.org/stable/api/mitmproxy/http.html#HTTPFlow).
+
+The `callback` addon is a Complement-Crypto specific addon which calls a client provided URL
+mid-flow, with a JSON object containing information about the HTTP flow. The caller can then
+return another JSON object which can modify the response in some way.
+
+Available configuration options are optional:
+ - `callback_request_url`: the URL to send outbound requests to. This allows callbacks to intercept
+   requests BEFORE they reach the server.
+ - `callback_response_url`: the URL to send inbound responses to. This allows callbacks to modify
+   response content.
+ - `filter`: the [mitmproxy filter](https://docs.mitmproxy.org/stable/concepts-filters/) to apply. If unset, ALL requests are eligible to go to the callback
+   server.
+
+To use this with the controller API, you would send an HTTP request like this:
+```js
+{
+  "options": {
+    "callback": {
+      "callback_response_url": "http://host.docker.internal:445566/response"
+    }
+  }
+}
+```
+
+#### `callback_request_url`
+
+mitmproxy will POST to `callback_request_url` with the following JSON object:
+```js
+{
+   method: "GET|PUT|...",
+   access_token: "syt_11...",
+   url: "http://hs1/_matrix/client/...",
+   request_body: { some json object or null if no body },
+}
+```
+The callback server can then either return an empty object or the following object (all fields are required):
+```js
+{
+   respond_status_code: 200,
+   respond_body: { "some": "json_object" }
+}
+```
+If an empty object is returned, mitmproxy will forward the request unaltered to the server. If the above object (with all fields set) is returned, mitmproxy will send that response _immediately_ and **will not send the request to the server**. This can be used to block HTTP requests.
+
+
+#### `callback_response_url`
+Similarly, mitmproxy will POST to `callback_response_url` with the following JSON object:
+```js
+{
+   method: "GET|PUT|...",
+   access_token: "syt_11...",
+   url: "http://hs1/_matrix/client/...",
+   request_body: { some json object or null if no body },
+   // note these are new fields because the request was sent to the HS and a response returned from it
+   response_body: { some json object },
+   response_code: 200,
+}
+```
+The callback server can then return optional keys to replace parts of the response.
+The values returned here will be returned to the Matrix client:
+```js
+{
+   respond_status_code: 200,
+   respond_body: { "some": "json_object" }
+}
+```
+These keys are optional. If neither are specified, the response is sent unaltered to
+the Matrix client. If the body is set but the status code is not, only the body is
+modified and the status code is left unaltered and vice versa.

tests/mitmproxy_addons/__init__.py

@@ -11,12 +11,10 @@ def install(package):
 install("aiohttp")
 
 from callback import Callback
-from status_code import StatusCode
 from controller import MITM_DOMAIN_NAME, app
 
 addons = [
     asgiapp.WSGIApp(app, MITM_DOMAIN_NAME, 80), # requests to this host will be routed to the flask app
-    StatusCode(),
     Callback(),
 ]
 # testcontainers will look for this log line

tests/mitmproxy_addons/callback.py

@@ -10,35 +10,7 @@
 from urllib.error import HTTPError, URLError
 from datetime import datetime
 
-# Callback will intercept a request and/or response and send a POST request to the provided url, with
-# the following JSON object. Supports filters: https://docs.mitmproxy.org/stable/concepts-filters/
-# {
-#   method: "GET|PUT|...",
-#   access_token: "syt_11...",
-#   url: "http://hs1/_matrix/client/...",
-#   request_body: { some json object or null if no body },
-#   response_body: { some json object },
-#   response_code: 200,
-# }
-# The response to this request can control what gets returned to the client. The response object:
-# {
-#    "respond_status_code": 200,
-#    "respond_body": { "some": "json_object" }
-# }
-# If {} is sent back, the response is not modified. Likewise, if `respond_body` is set but
-# `respond_status_code` is not, only the response body is modified, not the status code, and vice versa.
-#
-# To use this addon, configure it with these fields:
-#  - callback_request_url: the URL to send outbound requests to. This allows callbacks to intercept
-#                          requests BEFORE they reach the server. The request/response struct in this
-#                          callback is the same as `callback_response_url`, except `response_body`
-#                          and `response_code` will be missing as the request hasn't been processed
-#                          yet. To block the request from reaching the server, the callback needs to
-#                          provide both `respond_status_code` and `respond_body`.
-#  - callback_response_url: the URL to send inbound responses to. This allows callbacks to modify
-#                           response content.
-#  - filter: the mitmproxy filter to apply. If unset, ALL requests are eligible to go to the callback
-#            server.
+# See README.md for information about this addon
 class Callback:
     def __init__(self):
         self.reset()
@@ -149,6 +121,7 @@ async def send_callback(self, flow, url: str, body: dict):
                     # For responses: fields are optional but the default case is always specified. 
                     respond_status_code = test_response_body.get("respond_status_code", body.get("response_code"))
                     respond_body = test_response_body.get("respond_body", body.get("response_body"))
+                    print(f'{datetime.now().strftime("%H:%M:%S.%f")} callback for {flow.request.url} returning custom response: HTTP {respond_status_code} {json.dumps(respond_body)}')
                     flow.response = Response.make(
                         respond_status_code, json.dumps(respond_body),
                         headers={