Implement Context interceptor for ServeHTTP

[Zambiorix]

I need to adjust the Context when using ServeHTTP, when not feasible with normal middleware.

Consider the following scenario.

When using WebSockets, I want to convert a received message to a http.Request. Together with a custom ResponseWriter the request is handled by echo's ServeHTTP. But I want to adjust the Context and add my Socket reference to it.

There is no way to access the Context, except via middleware, but when handling middleware I have no actual access the the websocket anymore.

The easiest way to do this without interfering with other functionality or adding new functions, is checking if ResponseWriter conforms to a specific interface (ServeHTTPContextInterceptor) , if so, execute the interceptor and continue handling with the adjusted context.

Now a custom ResponseWriter can implement this interface and return a custom Context before handle and cleanup aftwards.

func (responseWriter *responseWriter) InterceptContext(ctx Echo.Context, handle func(ctx Echo.Context)) {

    handle(ctx)
}

Implemented and running in production at scale.

Another implemented usecase:

I have an Amazon AWS API Gateway routing all traffic to one Lambda Function, written in golang.

On receiving the call, I create a http.Request and ResponseWriter and perform ServeHTTP on echo, to handle the call.

Now I can create a Custom Context that hold various extra details about the request from API gateway.

[aldas]

Hi, could you build some working small example. I do not understand

There is no way to access the Context, except via middleware, but when handling middleware I have no actual access the the websocket anymore.

Why would this pre-middleware not work? You can add things to context.Context by replacing it in request or use methodecho.Context.Set.

package main

import (
    "context"
    "errors"
    "github.com/labstack/echo/v4"
    "log"
    "net/http"
)

type securityContext struct { // you can add methods to this object if you like
    UserID int64
}

func main() {
    e := echo.New()

    e.Pre(func(next echo.HandlerFunc) echo.HandlerFunc {
        return func(c echo.Context) error {
            sc := securityContext{
                UserID: 1,
            }

            req := c.Request()
            ctx := context.WithValue(req.Context(), "keyFromContextContext", sc)
            c.SetRequest(req.WithContext(ctx))

            // or
            c.Set("keyFromEchoContext", sc)

            return next(c)
        }
    })

    e.GET("/", func(c echo.Context) error {

        sc := (c.Request().Context().Value("keyFromContextContext")).(securityContext)
        sc2 := (c.Get("keyFromEchoContext")).(securityContext)

        return c.JSON(http.StatusOK, map[string]int64{
            "keyFromContextContext": sc.UserID,
            "keyFromEchoContext":    sc2.UserID,
        })
    })

    if err := e.Start(":8080"); err != nil && !errors.Is(err, http.ErrServerClosed) {
        log.Fatal(err)
    }
}

implementing echo.Context interface is something we do not encourage - the amount of methods you need to implement is a lot of unnecessary work. .Set method or your context.Conctext can return any object and it comes to almost same amount of code (get the value + cast it to type)

[Zambiorix]

I am not starting an echo server (I don't need one). I'm manually creating a http.Request and handle that via a ServeHTTP directly.

When using middleware, i have no access anymore to the data (I want to add to the context) at the time of ServeHTTP call.

Also, I,m not creating a completely new context, I'm just creating one inheriting the old, just to add some extra data.

Anyhow, I have several usecases for this concept that I could not solve with Pre or other middleware.

I'll see if I can make a quick test

[aldas]

but you do not need to start HTTP server.

package main

import (
    "context"
    "github.com/labstack/echo/v4"
    "net/http"
    "net/http/httptest"
)

type securityContext struct {
    UserID int64
}

func main() {
    e := echo.New()

    e.Pre(func(next echo.HandlerFunc) echo.HandlerFunc {
        return func(c echo.Context) error {
            sc := securityContext{
                UserID: 1,
            }

            req := c.Request()
            ctx := context.WithValue(req.Context(), "keyFromContextContext", sc)
            c.SetRequest(req.WithContext(ctx))

            // or
            c.Set("keyFromEchoContext", sc)

            return next(c)
        }
    })

    e.GET("/", func(c echo.Context) error {

        sc := (c.Request().Context().Value("keyFromContextContext")).(securityContext)
        sc2 := (c.Get("keyFromEchoContext")).(securityContext)

        return c.JSON(http.StatusOK, map[string]int64{
            "keyFromContextContext": sc.UserID,
            "keyFromEchoContext":    sc2.UserID,
        })
    })

    // somewhere (lamba stuff) serve the request and response with e.ServeHTTP

    // in this example we use httptest package
    req := httptest.NewRequest(http.MethodGet, "/", nil)
    rec := httptest.NewRecorder()

    e.ServeHTTP(rec, req) // this still executes that middleware and add values to context that can be accessed in handlers/middlewares

}

[aldas]

I think I need to see example. Premiddlewares and setting values is executed before this part https://github.com/labstack/echo/blob/f7d9f5142e9afa837c95b81a8bdc679ad807557e/echo.go#L665 is executed. so you can change pretty much everything in premiddleware as it is run before routing and "ordinary" middlewares

[Zambiorix]

1) I have 1000's events occurring per second 2) each event has unique access to an object relevant for that event 3) event creates request and calls ServeHTTP 4) somehow i need to attach unique event object to context when calling ServeHTTP 5) the place where I have to set the middleware has no access to this event data 6) I cannot possibly have to set middleware in my event function ...

A solution would be that in the event function I set a header with a token, push my object in some storage and fetch it in the middlware and then attach it to the context for use further down the chain.

But this is too costly.

Accessing the context in the ServeHTTP call before the chain is handled is exactly what I need, I believe.

Example on the way

[Zambiorix]

Test case added and functions as expected. Also, the cost is only 1 type assertion

[Zambiorix]

solved a small bug in the test

[Zambiorix]

I can rename the InterceptContext function in the responseWriter into EchoContextIntercept to make sure it never collides with anything else out there...

[aldas]

does the "value"/"socket reference" in context need to be accessed in that custom ResponseWriter and/or in handler? I am trying to isolate/pinpoint the requirements here.

1. if the need is to access value only in ResponseWriter at it to the custom rw.
2. if the need is to access value only in Handler set it to the context with echo.Context.Set.
3. if the need is to access value both in Handler and in ResponseWriter methods

you can create RW containing that holder and you can access it from handler and inside RW methods

package main_test

import (
    "bytes"
    "context"
    "github.com/labstack/echo/v4"
    "github.com/stretchr/testify/assert"
    "net/http"
    "testing"
)

type Holder struct {
    Socket string
}

type serveHTTPContextInterceptorRW struct {
    header     http.Header
    buffer     bytes.Buffer
    statusCode int

    holder *Holder
}

func (rw *serveHTTPContextInterceptorRW) Header() http.Header {
    if rw.header == nil {
        rw.header = make(http.Header)
    }
    return rw.header
}

func (rw *serveHTTPContextInterceptorRW) Write(data []byte) (int, error) {
    return rw.buffer.Write(data)
}

func (rw *serveHTTPContextInterceptorRW) WriteHeader(statusCode int) {
    rw.statusCode = statusCode
}

func TestEcho_RWContainsTheHolder(t *testing.T) {
    e := echo.New()

    e.GET("/foo", func(c echo.Context) error {
        if h, ok := c.Response().Writer.(*serveHTTPContextInterceptorRW); ok {
            return c.String(http.StatusOK, h.holder.Socket)
        }
        return c.NoContent(http.StatusBadRequest)
    })

    req, _ := http.NewRequest(http.MethodGet, "/foo", nil)
    //rw.intercept = func(ctx Context, handle func(ctx Context)) {
    //  handle(&context{ctx, "bar"})
    //}

    rw := &serveHTTPContextInterceptorRW{
        holder: &Holder{
            "bar",
        },
    }
    e.ServeHTTP(rw, req)

    assert.Equal(t, http.StatusOK, rw.statusCode)
    assert.Equal(t, "bar", rw.buffer.String())
}

or basically same but set the holder into context.Context and get the value from context in handler

func TestEcho_HolderIsInTheRWAndContext(t *testing.T) {
    e := echo.New()

    e.GET("/foo", func(c echo.Context) error {
        if h, ok := c.Request().Context().Value("myHolder").(*Holder); ok {
            return c.String(http.StatusOK, h.Socket)
        }
        return c.NoContent(http.StatusBadRequest)
    })

    holder := &Holder{
        "bar",
    }
    //rw.intercept = func(ctx Context, handle func(ctx Context)) {
    //  handle(&context{ctx, "bar"})
    //}

    req, _ := http.NewRequest(http.MethodGet, "/foo", nil)
    req = req.WithContext(context.WithValue(req.Context(), "myHolder", holder))
    rw := &serveHTTPContextInterceptorRW{
        holder: holder,
    }
    e.ServeHTTP(rw, req)

    assert.Equal(t, http.StatusOK, rw.statusCode)
    assert.Equal(t, "bar", rw.buffer.String())
}

and if that req.WithContext(context.WithValue(req.Context(), "myHolder", holder)) seems like costly operation consider this - http.Server uses it in multiple places to add things to the context when request is being served (before calling ServeHTTP)

[Zambiorix]

req.WithContext(context.WithValue(req.Context(), "myHolder", holder))

Looks like a good solution! Adding extra functions to ResponseWriter did look like a hack anyway ... Will close the PR

Thank you for your help!