This commit is contained in:
Saksham Arya 2024-11-22 07:44:49 +00:00 committed by GitHub
commit c8584677ce
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 29 additions and 2 deletions

View File

@ -71,6 +71,9 @@ type Context struct {
// This mutex protects Keys map.
mu sync.RWMutex
// This mutex protects headers map
hmu sync.RWMutex
// Keys is a key/value pair exclusively for the context of each request.
Keys map[string]any
@ -983,6 +986,8 @@ func (c *Context) Status(code int) {
// It writes a header in the response.
// If value == "", this method removes the header `c.Writer.Header().Del(key)`
func (c *Context) Header(key, value string) {
c.hmu.Lock()
defer c.hmu.Unlock()
if value == "" {
c.Writer.Header().Del(key)
return
@ -992,6 +997,8 @@ func (c *Context) Header(key, value string) {
// GetHeader returns value from request headers.
func (c *Context) GetHeader(key string) string {
c.hmu.RLock()
defer c.hmu.RUnlock()
return c.requestHeader(key)
}

View File

@ -166,7 +166,7 @@ func TestSaveUploadedFileWithPermission(t *testing.T) {
require.NoError(t, err)
mw.Close()
c, _ := CreateTestContext(httptest.NewRecorder())
c.Request, _ = http.NewRequest("POST", "/", buf)
c.Request, _ = http.NewRequest(http.MethodPost, "/", buf)
c.Request.Header.Set("Content-Type", mw.FormDataContentType())
f, err := c.FormFile("file")
require.NoError(t, err)
@ -187,7 +187,7 @@ func TestSaveUploadedFileWithPermissionFailed(t *testing.T) {
require.NoError(t, err)
mw.Close()
c, _ := CreateTestContext(httptest.NewRecorder())
c.Request, _ = http.NewRequest("POST", "/", buf)
c.Request, _ = http.NewRequest(http.MethodPost, "/", buf)
c.Request.Header.Set("Content-Type", mw.FormDataContentType())
f, err := c.FormFile("file")
require.NoError(t, err)
@ -3123,3 +3123,23 @@ func TestContextNext(t *testing.T) {
assert.True(t, exists)
assert.Equal(t, "value3", value)
}
func TestParallelHeaderWrite(t *testing.T) {
c, _ := CreateTestContext(httptest.NewRecorder())
wg := sync.WaitGroup{}
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < 1000; i++ {
c.Header("key", "value")
}
}()
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < 1000; i++ {
c.Header("key", "value")
}
}()
wg.Wait()
}