Migrate request context on mitmHandler

[hazcod]

Hi, so I have a TLS listener with my mps proxy.
During handshake, I generate a variable and need to pass this to my .OnRequest(...) function.
However, anything I set on req Context with the following function is not propagated?

func getInsertor(proxy *mps.HttpProxy) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req * http.Request) {
		logger := log.WithField("ip", req.RemoteAddr).WithField("host", req.Host)

		testID, err := pkg.ExtractTestID(req)
		if testID <= 0 || err != nil {
			logger.WithError(err).Error("could not extract testid")
			_, _ = rw.Write([]byte("invalid credentials provided"))
			rw.WriteHeader(http.StatusForbidden)
			return
		}
		proxy.ServeHTTP(rw, req.WithContext(context.WithValue(req.Context(), "testid", testID)))
	})
}

	proxyInstance := mps.NewHttpProxy()
	mitmHandler := mps.NewMitmHandler()

	mitmHandler.OnRequest().DoFunc(func(req *http.Request, ctx *mps.Context) (*http.Request, *http.Response){
		testID, ok := req.Context().Value("testid").(uint64) // This will always be nil/zero

[hazcod]

Same with ctx.Request.Context() or ctx.Context

[hazcod]

And strangely enough, any headers set in getInsertor does not seem to propagate to the OnRequest.
Edit: I guess this is because getInsertor only occurs on the CONNECT call, and does not modify the actual request embedded within.

[telanflow]

You should instantiate mitmHandler as follows:

proxyInstance := mps.NewHttpProxy()

// Use the same mps.Context instance
mitmHandler := mps.NewMitmHandlerWithContext(proxyInstance.Ctx)

// Process the CONNECT request
proxyInstance.HandleConnect = mitmHandler

Because middleware is implemented by mps.Context, Filter is based middleware

As middleware is implemented by mps.Context, middleware and filter won't work if you don't use the same mps.Context, which is a bit bad, so I wanted to find a better way to middleware.

[hazcod]

Ahh. not quite sure how I missed that, thanks @telanflow !
Loving the rewrite.

[hazcod]

@telanflow can you elaborate on which context I should set/get the value to get it to pass from CONNECT to the followup GET or POST request?
It doesn't seem to pass on either req.Context or ctx.Context.

I am using mitmHandler.OnRequest().DoFunc, but same with proxyInstance.OnRequest().DoFunc.

		if req.Method == http.MethodConnect {
			testID, err = pkg.ExtractTestID(req)
		} else {
		// if not, we really need a testID to be present in the context to continue
			var ok bool
			testID, ok = ctx.Context.Value("testid").(uint64)
			if !ok { err = errors.Errorf("invalid testid in context: %+v", req.Context().Value("testid")) }
		}

		if testID <= 0 || err != nil {
			logger.WithError(err).Error("could not extract testid")
			return req, newHttpResponse(req, http.StatusForbidden, "Forbidden", "invalid credentials provided")
		}

		if req.Method == http.MethodConnect {
			ctx.Context = context.WithValue(ctx.Context, "testid", testID)
			return req.WithContext(context.WithValue(req.Context(), "testid", testID)), nil
		}

[hazcod]

I tried

			req = req.WithContext(context.WithValue(req.Context(), "testid", testID))
			ctx = ctx.WithRequest(req)
			ctx.Context = context.WithValue(ctx.Context, "testid", testID)
			return req, nil

[telanflow]

I'm sorry, i went out to play on China National Day.

you can do that:

// set value
ctx.Context = context.WithValue(ctx.Context, "testid", testID)

// get value
ctx.Context.Value("testId")

Or

// set value
req.WithContext(context.WithValue(req.Context(), "testid", testID))

// Get value
req.Context().Value("testId")

[hazcod]

@telanflow is it possible it won't propagate to the actual requests because the GET request is inside the CONNECT tunnel, hence a different proxy context perhaps?

e.g.

CONNECT website.com:443
// setting context

GET /foo (in CONNECT tunnel)
// getting context -> nil

I am doing it as you stipulated:

			ctx.Context = context.WithValue(ctx.Context, "testid", testID)

[telanflow]

yes, It's best to set it above the request Context

[hazcod]

@telanflow can you perhaps elaborate? Not sure what you mean;

	proxyInstance.OnRequest().DoFunc(func(req *http.Request, ctx *mps.Context) (*http.Request, *http.Response) {
		logger := log.WithField("ip", req.RemoteAddr)

		var testID uint64
		var err error

		// if this is a CONNECT, scrape the testID from the request TLS metadata
		if req.Method == http.MethodConnect {
			testID, err = pkg.ExtractTestID(req)
		} else {
		// if not, we really need a testID to be present in the context to continue
			var ok bool
			testID, ok = ctx.Request.Context().Value("testid").(uint64)
			if !ok { err = errors.Errorf("invalid testid in context: %+v", req.Context().Value("testid")) }
		}

		if testID <= 0 || err != nil {
			logger.WithError(err).Error("could not extract testid")
			return req, newHttpResponse(req, http.StatusForbidden, "Forbidden", "invalid credentials provided")
		}

		// if this is a request to our proxy, set context and pass for the subsequent real requests to work
		if req.Method == http.MethodConnect {
			req = req.WithContext(context.WithValue(req.Context(), "testid", testID))
			ctx = ctx.WithRequest(req)
			ctx.Context = context.WithValue(ctx.Context, "testid", testID)
			return req, nil
		}
 ....
}

[telanflow]

Each request spawns a new Handler for the execution middleware (MITMHandler, TunnelHandler)

ctx := mitm.Ctx.WithRequest(req) 
resp, err := ctx.Next(req)

eg.

mps/mitm_handler.go

Lines 94 to 108 in 00583b5

	// Standard net/http function. You can use it alone
	func (mitm *MitmHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// execution middleware
	ctx := mitm.Ctx.WithRequest(req)
	resp, err := ctx.Next(req)
	if err != nil && err != MethodNotSupportErr {
	if resp != nil {
	copyHeaders(rw.Header(), resp.Header, mitm.Ctx.KeepDestinationHeaders)
	rw.WriteHeader(resp.StatusCode)
	buf := mitm.buffer().Get()
	_, err = io.CopyBuffer(rw, resp.Body, buf)
	mitm.buffer().Put(buf)
	}
	return
	}

mps/tunnel_handler.go

Lines 44 to 58 in 00583b5

	// Standard net/http function. You can use it alone
	func (tunnel *TunnelHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// execution middleware
	ctx := tunnel.Ctx.WithRequest(req)
	resp, err := ctx.Next(req)
	if err != nil && err != MethodNotSupportErr {
	if resp != nil {
	copyHeaders(rw.Header(), resp.Header, tunnel.Ctx.KeepDestinationHeaders)
	rw.WriteHeader(resp.StatusCode)
	buf := tunnel.buffer().Get()
	_, err = io.CopyBuffer(rw, resp.Body, buf)
	tunnel.buffer().Put(buf)
	}
	return
	}

[hazcod]

But proxy context should pass on between those handlers, right?