Bug in middleware or expected?

[spa5k]

I've been trying to make a session middleware for huma, based on SCS, at https://github.com/spa5k/huma-scs, but the problem I'm having is that it's not considering the cookies/anything that is set after the next(ctx) is called

// LoadAndSave provides middleware which automatically loads and saves session
// data for the current request, and communicates the session token to and from
// the client in a cookie.
func (s *SessionManager) LoadAndSave(ctx huma.Context, next func(huma.Context)) {
	ctx.SetHeader("Vary", "Cookie")
	var token string
	cookie, err := huma.ReadCookie(ctx, s.Cookie.Name)

	if err == nil {
		token = cookie.Value
	}

	newCtx, err := s.Load(ctx.Context(), token)
	if err != nil {
		s.ErrorFunc(ctx, err)
		return
	}

	ctx = huma.WithContext(ctx, newCtx)

	next(ctx)

	s.commitAndWriteSessionCookie(ctx)
}

I need a way to still make these headers available. If I put them above the next, they work, but I need them below for the logic to work.

It passes the test cases tho, the values come correctly for some reason, that is why it feels like a bug

[spa5k]

In here, both cache-control and the Set-Cookie header done in s.commitAndWriteSessionCookie, you can check the full code in the repo above

[spa5k]

Sample setup and routes to test the middleware -

humaConfig := huma.DefaultConfig("ReviewHQ API", "1")
api := humachi.New(r, humaConfig)
sessionManager := scs.New()
sessionManager.Lifetime = 1 * 365 * 24 * time.Hour
sessionManager.Cookie.Persist = true

api.UseMiddleware(sessionManager.LoadAndSave)

func sessionTestRoutes(api huma.API, sessionManager *scs.SessionManager) {
	huma.Register(api, huma.Operation{
		Method: http.MethodPut,
		Path:   "/put",
	}, func(ctx context.Context, input *struct{}) (*struct {
		Status  int          `json:"status"`
	}, error,
	) {
		sessionManager.Put(ctx, "foo", "bar")
		fmt.Println("writing session cookie")
		return &struct {
			Status  int          `json:"status"`
		}{Status: http.StatusOK}, nil
	})

	huma.Register(api, huma.Operation{
		Method: http.MethodGet,
		Path:   "/get",
	}, func(ctx context.Context, input *struct {
		Session string `cookie:"session"`
	}) (*struct {
		Status int    `json:"status"`
		Body   string `json:"body"`
	}, error,
	) {
		v := sessionManager.Get(ctx, "foo")
		if v == nil {
			return nil, huma.NewError(http.StatusInternalServerError, "foo does not exist in session")
		}

		return &struct {
			Status int    `json:"status"`
			Body   string `json:"body"`
		}{Status: http.StatusOK, Body: v.(string)}, nil
	})
}

[danielgtaylor]

@spa5k I'm not familiar with scs, but assuming you mean https://github.com/alexedwards/scs

the problem I'm having is that it's not considering the cookies/anything that is set after the next(ctx) is called

I think this would be the case with anything that tries to write headers (including cookies) to the response after data has already been written for the body. Even in the SCS example they have the session manager's Put(...) called within the handler itself and before any body is sent.

I think this is expected behavior, but I also understand what you are trying to do. You may have to do something like a defer in the handler which ensures the cookie is written before the body is serialized. You could make this generic by wrapping huma.Register with your own code that adds the defer, then using that function instead of huma.Register. Hope that helps!

[spa5k]

I've created a fork of SCS so that it works with huma - https://github.com/spa5k/huma-scs

The bug i'm asking about is here - https://github.com/spa5k/huma-scs/blob/5fd5277377d051a7163b513f9179ea31d56e9ba7/session.go#L150

The problem is that the way it works is that the user will add whatever he wants to the session context, like put/renew etc, and after it's done, and the request comes back from next(ctx), we will analyze what he has changed, and update the cookie that way. For now if you debug it, you will see that the huma.Context -> writer -> headers has the correct cookie related headers, but when the response is sent back, only the things above the next(ctx) are sent. So I mainly want to ask how we can still make it append the changes it make after next(ctx) to the response body.

[spa5k]

The huma.Register is just a sample on the routes, tbf I think we should have that feature just like the std library has, because manually adding the cookies and tracking everything will make it very verbose.

[spa5k]

Another reason why this is a bug is that it works correctly in the tests https://github.com/spa5k/huma-scs/blob/5fd5277377d051a7163b513f9179ea31d56e9ba7/session_test.go#L144

[danielgtaylor]

@spa5k yes this works in tests because the response writer is mutable after the body has been written, but in a real service that's not the case, you can't write headers after the body is sent (ignoring trailers for now). I'm open to ideas for how to address this, but I'm like 99% sure you'll get the same behavior without Huma if a handler method writes a body and the middleware tries to set a header afterward.

[spa5k]

hmm acceptable, so how you think we should be handling session management in huma, do you mind creating a basic plugin that adds/removes from the session and manages the cookies, then I will update it so that we can use scs stores etc, and maybe move it in huma organisation or something

[danielgtaylor]

@spa5k I think you could actually use a response transformer for this purpose as it gives you access to the Huma context before writing any response body. It should let you set headers as needed for both responses with and without a body. This chart shows where it fits into the request/response flow: https://huma.rocks/features/operations/#request-flow.

You would set up your transformer on the API config at https://github.com/danielgtaylor/huma/blob/main/api.go#L198-L199 so that it runs on every response. So you would use a middleware to set up the session, then a transformer to write the headers if needed. Hope that makes sense!

[spa5k]

Sure I'll take a look at it. If it works, then great, otherwise I'll keep spamming you until we get a session Middleware 😏