diff --git a/header.go b/header.go index 382e78b..3dee640 100644 --- a/header.go +++ b/header.go @@ -778,6 +778,8 @@ func (h *RequestHeader) VisitAllCookie(f func(key, value []byte)) { // // f must not retain references to key and/or value after returning. // Copy key and/or value contents before returning if you need retaining them. +// +// To get the headers in order they were received use VisitAllInOrder. func (h *RequestHeader) VisitAll(f func(key, value []byte)) { h.parseRawHeaders() host := h.Host() @@ -807,6 +809,25 @@ func (h *RequestHeader) VisitAll(f func(key, value []byte)) { } } +// VisitAllInOrder calls f for each header in the order they were received. +// +// f must not retain references to key and/or value after returning. +// Copy key and/or value contents before returning if you need retaining them. +// +// This function is slightly slower than VisitAll because it has to reparse the +// raw headers to get the order. +func (h *RequestHeader) VisitAllInOrder(f func(key, value []byte)) { + h.parseRawHeaders() + var s headerScanner + s.b = h.rawHeaders + s.disableNormalizing = h.disableNormalizing + for s.next() { + if len(s.key) > 0 { + f(s.key, s.value) + } + } +} + // Del deletes header with the given key. func (h *ResponseHeader) Del(key string) { k := getHeaderKeyBytes(&h.bufKV, key, h.disableNormalizing) diff --git a/header_test.go b/header_test.go index 57770bb..e0f567f 100644 --- a/header_test.go +++ b/header_test.go @@ -1277,6 +1277,47 @@ func TestRequestHeaderVisitAll(t *testing.T) { } } +func TestResponseHeaderVisitAllInOrder(t *testing.T) { + var h RequestHeader + + r := bytes.NewBufferString("GET / HTTP/1.1\r\nContent-Type: aa\r\nCookie: a=b\r\nHost: example.com\r\nUser-Agent: xxx\r\n\r\n") + br := bufio.NewReader(r) + if err := h.Read(br); err != nil { + t.Fatalf("Unepxected error: %s", err) + } + + if h.Len() != 4 { + t.Fatalf("Unexpected number of headers: %d. Expected 4", h.Len()) + } + + order := []string{ + "Content-Type", + "Cookie", + "Host", + "User-Agent", + } + values := []string{ + "aa", + "a=b", + "example.com", + "xxx", + } + + h.VisitAllInOrder(func(key, value []byte) { + if len(order) == 0 { + t.Fatalf("no more headers expected, got %q", key) + } + if order[0] != string(key) { + t.Fatalf("expected header %q got %q", order[0], key) + } + if values[0] != string(value) { + t.Fatalf("expected header value %q got %q", values[0], value) + } + order = order[1:] + values = values[1:] + }) +} + func TestResponseHeaderCookie(t *testing.T) { var h ResponseHeader var c Cookie