goravel · hwbrzzl · Dec 15, 2024 · Dec 15, 2024 · Dec 15, 2024 · Dec 15, 2024
@@ -9,6 +9,21 @@ import (
 	"github.com/valyala/fasthttp"
 )
 
+type contextKeyType struct{}
+type sessionKeyType struct{}
+type userContextKeyType struct{}
+
+var (
+	contextKey          = contextKeyType{}
+	sessionKey          = sessionKeyType{}
+	userContextKey      = userContextKeyType{}
+	internalContextKeys = []any{
+		contextKey,
+		sessionKey,
+		userContextKey,
+	}
+)
+
 func Background() http.Context {
 	app := fiber.New()
 	httpCtx := app.AcquireCtx(&fasthttp.RequestCtx{})
@@ -38,16 +53,40 @@ func (c *Context) Response() http.ContextResponse {
 }
 
 func (c *Context) WithValue(key any, value any) {
-	ctx := context.WithValue(c.instance.UserContext(), key, value)
+	// Not store the value in the context directly, because we want to return the value map when calling `Context()`.
+	values := c.getGoravelContextValues()
+	if values == nil {
+		values = make(map[any]any)
+	}
+	values[key] = value
+	ctx := context.WithValue(c.instance.UserContext(), contextKey, values)
 	c.instance.SetUserContext(ctx)
 }
 
 func (c *Context) WithContext(ctx context.Context) {
+	// We want to return the original context back when calling `Context()`, so we need to store it.
+	ctx = context.WithValue(ctx, userContextKey, ctx)
+	ctx = context.WithValue(ctx, contextKey, c.getGoravelContextValues())
 	c.instance.SetUserContext(ctx)
 }
 
 func (c *Context) Context() context.Context {
-	return c.instance.UserContext()
+	ctx := c.getUserContext()
+	values := c.getGoravelContextValues()
+	for key, value := range values {
+		skip := false
+		for _, internalContextKey := range internalContextKeys {
+			if key == internalContextKey {
+				skip = true
+			}
+		}
+
+		if !skip {
+			ctx = context.WithValue(ctx, key, value)
+		}
+	}
+
+	return ctx
 }
 
 func (c *Context) Deadline() (deadline time.Time, ok bool) {
@@ -63,9 +102,32 @@ func (c *Context) Err() error {
 }
 
 func (c *Context) Value(key any) any {
+	values := c.getGoravelContextValues()
+	if value, exist := values[key]; exist {
+		return value
+	}
+
 	return c.instance.UserContext().Value(key)
 }
 
 func (c *Context) Instance() *fiber.Ctx {
 	return c.instance
 }
+
+func (c *Context) getGoravelContextValues() map[any]any {
+	value := c.instance.UserContext().Value(contextKey)
+	if goravelCtxVal, ok := value.(map[any]any); ok {
+		return goravelCtxVal
+	}
+
+	return make(map[any]any)
+}
+
+func (c *Context) getUserContext() context.Context {
+	ctx, exist := c.instance.UserContext().Value(userContextKey).(context.Context)
+	if !exist {
+		ctx = context.Background()
+	}
+
+	return ctx
+}
@@ -139,7 +139,7 @@ func (r *ContextRequest) Host() string {
 }
 
 func (r *ContextRequest) HasSession() bool {
-	_, ok := r.ctx.Value("session").(contractsession.Session)
+	_, ok := r.ctx.Value(sessionKey).(contractsession.Session)
 	return ok
 }
 
@@ -392,15 +392,15 @@ func (r *ContextRequest) RouteInt64(key string) int64 {
 }
 
 func (r *ContextRequest) Session() contractsession.Session {
-	s, ok := r.ctx.Value("session").(contractsession.Session)
+	s, ok := r.ctx.Value(sessionKey).(contractsession.Session)
 	if !ok {
 		return nil
 	}
 	return s
 }
 
 func (r *ContextRequest) SetSession(session contractsession.Session) contractshttp.ContextRequest {
-	r.ctx.WithValue("session", session)
+	r.ctx.WithValue(sessionKey, session)
 
 	return r
 }

@@ -1,33 +1,48 @@
 package fiber
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 )
 
 func TestContext(t *testing.T) {
-	type customKeyType struct{}
-	var customKey customKeyType
+	type customType struct{}
+	var customTypeKey customType
 
 	httpCtx := Background()
 	httpCtx.WithValue("Hello", "world")
 	httpCtx.WithValue("Hi", "Goravel")
-	httpCtx.WithValue(customKey, "halo")
+	httpCtx.WithValue(customTypeKey, "halo")
+	httpCtx.WithValue(sessionKey, "session")
+
+	userContext := context.Background()
+	//nolint:all
+	userContext = context.WithValue(userContext, "user_a", "b")
+	httpCtx.WithContext(userContext)
+
 	httpCtx.WithValue(1, "one")
 	httpCtx.WithValue(2.2, "two point two")
 
 	assert.Equal(t, "world", httpCtx.Value("Hello"))
 	assert.Equal(t, "Goravel", httpCtx.Value("Hi"))
-	assert.Equal(t, "halo", httpCtx.Value(customKey))
+	assert.Equal(t, "halo", httpCtx.Value(customTypeKey))
 	assert.Equal(t, "one", httpCtx.Value(1))
 	assert.Equal(t, "two point two", httpCtx.Value(2.2))
+	assert.Equal(t, "session", httpCtx.Value(sessionKey))
+	assert.NotNil(t, httpCtx.Value(contextKey))
+
+	// The value of UserContext can't be covered.
+	assert.Equal(t, "b", httpCtx.Value("user_a"))
 
 	ctx := httpCtx.Context()
 	assert.Equal(t, "world", ctx.Value("Hello"))
 	assert.Equal(t, "Goravel", ctx.Value("Hi"))
-	assert.Equal(t, "halo", ctx.Value(customKey))
+	assert.Equal(t, "halo", ctx.Value(customTypeKey))
+	assert.Equal(t, "b", ctx.Value("user_a"))
 	assert.Equal(t, "one", ctx.Value(1))
 	assert.Equal(t, "two point two", ctx.Value(2.2))
+	assert.Nil(t, ctx.Value(sessionKey))
+	assert.Nil(t, ctx.Value(contextKey))
 }
-