Skip to content

Commit

Permalink
fix Request.Context() checks (#3512)
Browse files Browse the repository at this point in the history
Co-authored-by: Bence Vidosits <bence.vidosits1@ibm.com>
  • Loading branch information
Bence Vidosits and Bence Vidosits authored May 29, 2023
1 parent 2d4bbec commit bb1fc2e
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 4 deletions.
15 changes: 11 additions & 4 deletions context.go
Original file line number Diff line number Diff line change
Expand Up @@ -1180,25 +1180,32 @@ func (c *Context) SetAccepted(formats ...string) {
/***** GOLANG.ORG/X/NET/CONTEXT *****/
/************************************/

// hasRequestContext returns whether c.Request has Context and fallback.
func (c *Context) hasRequestContext() bool {
hasFallback := c.engine != nil && c.engine.ContextWithFallback
hasRequestContext := c.Request != nil && c.Request.Context() != nil
return hasFallback && hasRequestContext
}

// Deadline returns that there is no deadline (ok==false) when c.Request has no Context.
func (c *Context) Deadline() (deadline time.Time, ok bool) {
if !c.engine.ContextWithFallback || c.Request == nil || c.Request.Context() == nil {
if !c.hasRequestContext() {
return
}
return c.Request.Context().Deadline()
}

// Done returns nil (chan which will wait forever) when c.Request has no Context.
func (c *Context) Done() <-chan struct{} {
if !c.engine.ContextWithFallback || c.Request == nil || c.Request.Context() == nil {
if !c.hasRequestContext() {
return nil
}
return c.Request.Context().Done()
}

// Err returns nil when c.Request has no Context.
func (c *Context) Err() error {
if !c.engine.ContextWithFallback || c.Request == nil || c.Request.Context() == nil {
if !c.hasRequestContext() {
return nil
}
return c.Request.Context().Err()
Expand All @@ -1219,7 +1226,7 @@ func (c *Context) Value(key any) any {
return val
}
}
if !c.engine.ContextWithFallback || c.Request == nil || c.Request.Context() == nil {
if !c.hasRequestContext() {
return nil
}
return c.Request.Context().Value(key)
Expand Down
18 changes: 18 additions & 0 deletions context_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2176,6 +2176,24 @@ func TestRemoteIPFail(t *testing.T) {
assert.False(t, trust)
}

func TestHasRequestContext(t *testing.T) {
c, _ := CreateTestContext(httptest.NewRecorder())
assert.False(t, c.hasRequestContext(), "no request, no fallback")
c.engine.ContextWithFallback = true
assert.False(t, c.hasRequestContext(), "no request, has fallback")
c.Request, _ = http.NewRequest(http.MethodGet, "/", nil)
assert.True(t, c.hasRequestContext(), "has request, has fallback")
c.Request, _ = http.NewRequestWithContext(nil, "", "", nil) //nolint:staticcheck
assert.False(t, c.hasRequestContext(), "has request with nil ctx, has fallback")
c.engine.ContextWithFallback = false
assert.False(t, c.hasRequestContext(), "has request, no fallback")

c = &Context{}
assert.False(t, c.hasRequestContext(), "no request, no engine")
c.Request, _ = http.NewRequest(http.MethodGet, "/", nil)
assert.False(t, c.hasRequestContext(), "has request, no engine")
}

func TestContextWithFallbackDeadlineFromRequestContext(t *testing.T) {
c, _ := CreateTestContext(httptest.NewRecorder())
// enable ContextWithFallback feature flag
Expand Down

0 comments on commit bb1fc2e

Please sign in to comment.