Skip to content

Commit

Permalink
SetHeader should replace entries with the same key (#1134)
Browse files Browse the repository at this point in the history
* fix: SetHeader should replace entries with the same key

* feat: AddHeader function that appends the headers list

* fix: comment format
  • Loading branch information
lihe07 authored Nov 25, 2024
1 parent 73907a8 commit 3025dde
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 0 deletions.
25 changes: 25 additions & 0 deletions hijack.go
Original file line number Diff line number Diff line change
Expand Up @@ -356,6 +356,31 @@ func (ctx *HijackResponse) Headers() http.Header {

// SetHeader of the payload via key-value pairs.
func (ctx *HijackResponse) SetHeader(pairs ...string) *HijackResponse {
headerIndex := make(map[string]int, len(ctx.payload.ResponseHeaders))
for i, header := range ctx.payload.ResponseHeaders {
headerIndex[header.Name] = i
}

for i := 0; i < len(pairs); i += 2 {
name := pairs[i]
value := pairs[i+1]

if idx, exists := headerIndex[name]; exists {
ctx.payload.ResponseHeaders[idx].Value = value
} else {
ctx.payload.ResponseHeaders = append(ctx.payload.ResponseHeaders, &proto.FetchHeaderEntry{
Name: name,
Value: value,
})
headerIndex[name] = len(ctx.payload.ResponseHeaders) - 1
}
}
return ctx
}

// AddHeader appends key-value pairs to the end of the response headers.
// Duplicate keys will be preserved.
func (ctx *HijackResponse) AddHeader(pairs ...string) *HijackResponse {
for i := 0; i < len(pairs); i += 2 {
ctx.payload.ResponseHeaders = append(ctx.payload.ResponseHeaders, &proto.FetchHeaderEntry{
Name: pairs[i],
Expand Down
2 changes: 2 additions & 0 deletions hijack_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,8 @@ func TestHijack(t *testing.T) {
g.Has(ctx.Response.Headers().Get("Content-Type"), "text/html; charset=utf-8")

// override response header
ctx.Response.AddHeader("Set-Cookie", "key=val1")
// This should override the previous one
ctx.Response.SetHeader("Set-Cookie", "key=val")

// override response body
Expand Down

0 comments on commit 3025dde

Please sign in to comment.