Merge pull request #120 from matrix-org/kegan/mitm-api-4

Remove old MITM API and use new format

internal/deploy/callback/callback_addon.go

@@ -162,12 +162,13 @@ func NewCallbackServer(t ct.TestLike, hostnameRunningComplement string) (*Callba
 
 // SendError returns a callback.Fn which returns the provided statusCode
 // along with a JSON error $count times, after which it lets the response
-// pass through. This is useful for testing retries.
+// pass through. This is useful for testing retries. If count=0, always send
+// an error response.
 func SendError(count uint32, statusCode int) Fn {
 	var seen atomic.Uint32
 	return func(d Data) *Response {
 		next := seen.Add(1)
-		if next > count {
+		if count > 0 && next > count {
 			return nil
 		}
 		return &Response{

internal/deploy/mitm/client.go

@@ -6,7 +6,6 @@ import (
 	"io"
 	"net/http"
 	"net/url"
-	"sync"
 	"testing"
 	"time"
 
@@ -40,10 +39,8 @@ func NewClient(proxyURL *url.URL, hostnameRunningComplement string) *Client {
 
 func (m *Client) Configure(t *testing.T) *Configuration {
 	return &Configuration{
-		t:        t,
-		pathCfgs: make(map[string]*MITMPathConfiguration),
-		mu:       &sync.Mutex{},
-		client:   m,
+		t:      t,
+		client: m,
 	}
 }
 

internal/deploy/mitm/configuration.go

@@ -1,14 +1,10 @@
 package mitm
 
 import (
-	"encoding/json"
 	"strings"
-	"sync"
-	"sync/atomic"
 	"testing"
 
 	"github.com/matrix-org/complement-crypto/internal/deploy/callback"
-	"github.com/matrix-org/complement/ct"
 	"github.com/matrix-org/complement/must"
 )
 
@@ -17,10 +13,8 @@ import (
 // Tests will typically build up this configuration by calling `Intercept` with the paths
 // they are interested in.
 type Configuration struct {
-	t        *testing.T
-	pathCfgs map[string]*MITMPathConfiguration
-	mu       *sync.Mutex
-	client   *Client
+	t      *testing.T
+	client *Client
 }
 
 // Filter represents a mitmproxy filter; see https://docs.mitmproxy.org/stable/concepts-filters/
@@ -133,142 +127,3 @@ func (c *Configuration) WithIntercept(opts InterceptOpts, inner func()) {
 	defer c.client.UnlockOptions(c.t, lockID)
 	inner()
 }
-
-func (c *Configuration) ForPath(partialPath string) *MITMPathConfiguration {
-	c.mu.Lock()
-	defer c.mu.Unlock()
-	p, ok := c.pathCfgs[partialPath]
-	if ok {
-		return p
-	}
-	p = &MITMPathConfiguration{
-		t:    c.t,
-		path: partialPath,
-	}
-	c.pathCfgs[partialPath] = p
-	return p
-}
-
-// Execute a mitm proxy configuration for the duration of `inner`.
-func (c *Configuration) Execute(inner func()) {
-	// The HTTP request to mitmproxy needs to look like:
-	//   {
-	//     $addon_name: {
-	//       $addon_values...
-	//     }
-	//   }
-	//
-	// The API shape of the add-ons are located inside the python files in tests/mitmproxy_addons
-	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 := callback.NewCallbackServer(c.t, c.client.hostnameRunningComplement)
-		must.NotError(c.t, "failed to start callback server", err)
-		defer cbServer.Close()
-
-		if pathConfig.listener != nil {
-			responseCallbackURL := cbServer.SetOnResponseCallback(c.t, pathConfig.listener)
-			callbackAddon["callback_response_url"] = responseCallbackURL
-		}
-		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 callback.Data) *callback.Response {
-				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 &callback.Response{
-					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, map[string]any{
-		"callback": callbackAddon,
-	})
-	defer c.client.UnlockOptions(c.t, lockID)
-	inner()
-
-}
-
-type MITMPathConfiguration struct {
-	t           *testing.T
-	path        string
-	accessToken string
-	method      string
-	listener    func(cd callback.Data) *callback.Response
-
-	blockCount      int
-	blockStatusCode int
-	blockRequest    *bool // nil means don't block
-}
-
-func (p *MITMPathConfiguration) filter() string {
-	// the filter is a python regexp
-	// "Regexes are Python-style" - https://docs.mitmproxy.org/stable/concepts-filters/
-	// re.escape() escapes very little:
-	// "Changed in version 3.7: Only characters that can have special meaning in a regular expression are escaped.
-	// As a result, '!', '"', '%', "'", ',', '/', ':', ';', '<', '=', '>', '@', and "`" are no longer escaped."
-	// https://docs.python.org/3/library/re.html#re.escape
-	//
-	// The majority of HTTP paths are just /foo/bar with % for path-encoding e.g @foo:bar=>%40foo%3Abar,
-	// so on balance we can probably just use the path directly.
-	var s strings.Builder
-	s.WriteString("~u .*" + p.path + ".*")
-	if p.method != "" {
-		s.WriteString(" ~m " + strings.ToUpper(p.method))
-	}
-	if p.accessToken != "" {
-		s.WriteString(" ~hq " + p.accessToken)
-	}
-	return s.String()
-}
-
-func (p *MITMPathConfiguration) Listen(cb func(cd callback.Data) *callback.Response) *MITMPathConfiguration {
-	p.listener = cb
-	return p
-}
-
-func (p *MITMPathConfiguration) AccessToken(accessToken string) *MITMPathConfiguration {
-	p.accessToken = accessToken
-	return p
-}
-
-func (p *MITMPathConfiguration) Method(method string) *MITMPathConfiguration {
-	p.method = method
-	return p
-}
-
-func (p *MITMPathConfiguration) BlockRequest(count, returnStatusCode int) *MITMPathConfiguration {
-	if p.blockRequest != nil {
-		// we can't express blocking requests and responses separately, it doesn't make sense.
-		ct.Fatalf(p.t, "BlockRequest or BlockResponse cannot be called multiple times for the same path")
-	}
-	p.blockCount = count
-	p.blockRequest = &boolTrue
-	p.blockStatusCode = returnStatusCode
-	return p
-}
-
-func (p *MITMPathConfiguration) BlockResponse(count, returnStatusCode int) *MITMPathConfiguration {
-	if p.blockRequest != nil {
-		ct.Fatalf(p.t, "BlockRequest or BlockResponse cannot be called multiple times for the same path")
-	}
-	p.blockCount = count
-	p.blockRequest = &boolFalse
-	p.blockStatusCode = returnStatusCode
-	return p
-}

tests/notification_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/matrix-org/complement-crypto/internal/api"
 	"github.com/matrix-org/complement-crypto/internal/cc"
 	"github.com/matrix-org/complement-crypto/internal/deploy/callback"
+	"github.com/matrix-org/complement-crypto/internal/deploy/mitm"
 	"github.com/matrix-org/complement/must"
 )
 
@@ -567,23 +568,26 @@ func TestMultiprocessDupeOTKUpload(t *testing.T) {
 	// artificially slow down the HTTP responses, such that we will potentially have 2 in-flight /keys/upload requests
 	// at once. If the NSE and main apps are talking to each other, they should be using the same key ID + key.
 	// If not... well, that's a bug because then the client will forget one of these keys.
-	mitmConfiguration := tc.Deployment.MITM().Configure(t)
-	mitmConfiguration.ForPath("/keys/upload").Listen(func(cd callback.Data) *callback.Response {
-		if cd.AccessToken != aliceAccessToken {
-			return nil // let bob upload OTKs
-		}
-		aliceUploadedNewKeys = true
-		if cd.ResponseCode != 200 {
-			// we rely on the homeserver checking and rejecting when the same key ID is used with
-			// different keys.
-			t.Errorf("/keys/upload returned an error, duplicate key upload? %+v => %v", cd, string(cd.ResponseBody))
-		}
-		// tarpit the response
-		t.Logf("tarpitting keys/upload response for 4 seconds")
-		time.Sleep(4 * time.Second)
-		return nil
-	})
-	mitmConfiguration.Execute(func() {
+	tc.Deployment.MITM().Configure(t).WithIntercept(mitm.InterceptOpts{
+		Filter: mitm.FilterParams{
+			PathContains: "/keys/upload",
+		},
+		ResponseCallback: func(cd callback.Data) *callback.Response {
+			if cd.AccessToken != aliceAccessToken {
+				return nil // let bob upload OTKs
+			}
+			aliceUploadedNewKeys = true
+			if cd.ResponseCode != 200 {
+				// we rely on the homeserver checking and rejecting when the same key ID is used with
+				// different keys.
+				t.Errorf("/keys/upload returned an error, duplicate key upload? %+v => %v", cd, string(cd.ResponseBody))
+			}
+			// tarpit the response
+			t.Logf("tarpitting keys/upload response for 4 seconds")
+			time.Sleep(4 * time.Second)
+			return nil
+		},
+	}, func() {
 		var eventID string
 		// Bob appears and sends a message, causing Bob to claim one of Alice's OTKs.
 		// The main app will see this in /sync and then try to upload another OTK, which we will tarpit.

tests/one_time_keys_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/matrix-org/complement-crypto/internal/api"
 	"github.com/matrix-org/complement-crypto/internal/cc"
 	"github.com/matrix-org/complement-crypto/internal/deploy/callback"
+	"github.com/matrix-org/complement-crypto/internal/deploy/mitm"
 	"github.com/matrix-org/complement/b"
 	"github.com/matrix-org/complement/client"
 	"github.com/matrix-org/complement/ct"
@@ -108,9 +109,14 @@ func TestFallbackKeyIsUsedIfOneTimeKeysRunOut(t *testing.T) {
 			var roomID string
 			var waiter api.Waiter
 			// Block all /keys/upload requests for Alice
-			mitmConfiguration := tc.Deployment.MITM().Configure(t)
-			mitmConfiguration.ForPath("/keys/upload").AccessToken(alice.CurrentAccessToken(t)).BlockRequest(0, http.StatusGatewayTimeout)
-			mitmConfiguration.Execute(func() {
+			tc.Deployment.MITM().Configure(t).WithIntercept(mitm.InterceptOpts{
+				Filter: mitm.FilterParams{
+					PathContains: "/keys/upload",
+					Method:       "POST",
+					AccessToken:  alice.CurrentAccessToken(t),
+				},
+				RequestCallback: callback.SendError(0, http.StatusGatewayTimeout),
+			}, func() {
 				// claim all OTKs
 				mustClaimOTKs(t, otkGobbler, tc.Alice, int(otkCount))
 
@@ -156,9 +162,13 @@ func TestFailedOneTimeKeyUploadRetries(t *testing.T) {
 		// 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
-		mitmConfiguration := tc.Deployment.MITM().Configure(t)
-		mitmConfiguration.ForPath("/keys/upload").Method("POST").BlockRequest(2, http.StatusGatewayTimeout)
-		mitmConfiguration.Execute(func() {
+		tc.Deployment.MITM().Configure(t).WithIntercept(mitm.InterceptOpts{
+			Filter: mitm.FilterParams{
+				PathContains: "/keys/upload",
+				Method:       "POST",
+			},
+			RequestCallback: callback.SendError(2, http.StatusGatewayTimeout),
+		}, func() {
 			tc.WithAliceSyncing(t, func(alice api.Client) {
 				tc.Bob.MustDo(t, "POST", []string{
 					"_matrix", "client", "v3", "keys", "claim",
@@ -204,16 +214,21 @@ func TestFailedKeysClaimRetries(t *testing.T) {
 			// make a room which will link the 2 users together when
 			roomID := tc.CreateNewEncryptedRoom(t, tc.Alice, cc.EncRoomOptions.PresetPublicChat())
 			// block /keys/claim and join the room, causing the Olm session to be created
-			mitmConfiguration := tc.Deployment.MITM().Configure(t)
-			mitmConfiguration.ForPath("/keys/claim").Method("POST").BlockRequest(2, http.StatusGatewayTimeout).Listen(func(cd callback.Data) *callback.Response {
-				t.Logf("%+v", cd)
-				if cd.ResponseCode == 200 {
-					waiter.Finish()
-					stopPoking.Store(true)
-				}
-				return nil
-			})
-			mitmConfiguration.Execute(func() {
+			tc.Deployment.MITM().Configure(t).WithIntercept(mitm.InterceptOpts{
+				Filter: mitm.FilterParams{
+					PathContains: "/keys/claim",
+					Method:       "POST",
+				},
+				RequestCallback: callback.SendError(2, http.StatusGatewayTimeout),
+				ResponseCallback: func(cd callback.Data) *callback.Response {
+					t.Logf("%+v", cd)
+					if cd.ResponseCode == 200 {
+						waiter.Finish()
+						stopPoking.Store(true)
+					}
+					return nil
+				},
+			}, func() {
 				// join the room. This should cause an Olm session to be made but it will fail as we cannot
 				// call /keys/claim. We should retry though.
 				tc.Bob.MustJoinRoom(t, roomID, []string{clientType.HS})