Skip to content

Commit

Permalink
fix: SetHeader should replace entries with the same key
Browse files Browse the repository at this point in the history
  • Loading branch information
lihe07 committed Oct 30, 2024
1 parent 0096e53 commit cdde9e1
Showing 1 changed file with 17 additions and 4 deletions.
21 changes: 17 additions & 4 deletions hijack.go
Original file line number Diff line number Diff line change
Expand Up @@ -356,11 +356,24 @@ 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 {
ctx.payload.ResponseHeaders = append(ctx.payload.ResponseHeaders, &proto.FetchHeaderEntry{
Name: pairs[i],
Value: pairs[i+1],
})
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
}
Expand Down

0 comments on commit cdde9e1

Please sign in to comment.