rewrite Cookie with KVBuffer

This commit is contained in:
Patricio Whittingslow
2026-07-29 13:18:17 -03:00
parent feecdb309e
commit af0f1b8cc6
4 changed files with 46 additions and 72 deletions
+28 -65
View File
@@ -7,68 +7,60 @@ import (
// Cookie implements cookie key-value parsing. Methods function similarly to eponymous [Header] methods. // Cookie implements cookie key-value parsing. Methods function similarly to eponymous [Header] methods.
// Cookie represents a single-line Cookie header value in a HTTP header, much like the standard library Cookie. // Cookie represents a single-line Cookie header value in a HTTP header, much like the standard library Cookie.
type Cookie struct { type Cookie struct {
buf []byte kv KVBuffer
kvs []argsKV // first key-value pair is the data Key/Value pair.
} }
// Reset functions very similarly to [Header.Reset]. Can be used for in-place cookie parsing. // Reset functions very similarly to [Header.Reset]. Can be used for in-place cookie parsing.
func (c *Cookie) Reset(buf []byte) { func (c *Cookie) Reset(buf []byte, capKV int) { c.kv.Reset(buf, capKV) }
if buf == nil {
buf = c.buf[:0] func (c *Cookie) valid() bool {
} return len(c.kv.kvs) > 0 && c.kv.kvs[0].key.len > 0
*c = Cookie{
buf: buf,
kvs: c.kvs[:0],
}
} }
// Name returns the first cookie key which is commonly referred to as the cookie's name. Returns nil if not found. // Name returns the first cookie key which is commonly referred to as the cookie's name. Returns nil if not found.
func (c *Cookie) Name() []byte { func (c *Cookie) Name() []byte {
if len(c.kvs) == 0 || c.kvs[0].key.len == 0 { if !c.valid() {
return nil return nil
} }
return tok2bytes(c.buf, c.kvs[0].key) return c.kv.AtKey(0)
} }
// Value returns the first cookie value associated with the name. Returns nil if not found. // Value returns the first cookie value associated with the name. Returns nil if not found.
func (c *Cookie) Value() []byte { func (c *Cookie) Value() []byte {
if len(c.kvs) == 0 || c.kvs[0].value.len == 0 { if !c.valid() {
return nil return nil
} }
return tok2bytes(c.buf, c.kvs[0].value) return c.kv.AtValue(0)
} }
// ParseBytes copies the argument bytes to the Cookie's underlying buffer and parses the cookie. // ParseBytes copies the argument bytes to the Cookie's underlying buffer and parses the cookie.
func (c *Cookie) ParseBytes(cookie []byte) error { func (c *Cookie) ParseBytes(cookie []byte) error {
c.Reset(nil) c.Reset(nil, 0)
c.buf = append(c.buf[:0], cookie...) c.kv.buf = append(c.kv.buf[:0], cookie...)
return c.Parse() return c.Parse()
} }
// CopyFrom makes a copy of the argument cookie to the receiver dst argument. No memory is shared between cookies. // CopyFrom makes a copy of the argument cookie to the receiver dst argument. No memory is shared between cookies.
func (dst *Cookie) CopyFrom(c Cookie) { func (dst *Cookie) CopyFrom(c Cookie) { dst.kv.CopyFrom(&c.kv) }
dst.buf = append(dst.buf[:0], c.buf...)
dst.kvs = append(dst.kvs[:0], c.kvs...)
}
// Parse parses the cookie's buffer in place. // Parse parses the cookie's buffer in place.
func (c *Cookie) Parse() error { func (c *Cookie) Parse() error {
if len(c.kvs) > 0 { if c.kv.Len() > 0 {
return errCookiesParsed return errCookiesParsed
} }
off := 0 off := 0
for { for {
k, v, n := parseCookie(c.buf[off:]) k, v, n := parseCookie(c.kv.buf[off:])
if n == 0 { if n == 0 {
break break
} }
c.kvs = append(c.kvs, argsKV{ if !c.kv.setInternal(k, v) {
key: bytes2tok(c.buf, k), return errOOM
value: bytes2tok(c.buf, v), }
})
off += n off += n
} }
if len(c.kvs) == 0 { if c.kv.Len() == 0 {
return errNoCookies return errNoCookies
} }
return nil return nil
@@ -76,44 +68,17 @@ func (c *Cookie) Parse() error {
// ForEach iterates over the cookie's key-value pairs, stopping on the first // ForEach iterates over the cookie's key-value pairs, stopping on the first
// error returned by cb and returning it. // error returned by cb and returning it.
func (c *Cookie) ForEach(cb func(key, value []byte) error) error { func (c *Cookie) ForEach(cb func(key, value []byte) bool) {
nc := len(c.kvs) c.kv.ForEach(cb)
for i := range nc {
kv := c.kvs[i]
key := tok2bytes(c.buf, kv.key)
value := tok2bytes(c.buf, kv.value)
err := cb(key, value)
if err != nil {
return err
}
}
return nil
} }
// Get gets a cookie's value from its key. Use HasValueOrKey to check if a key or single-valued cookie is present in the cookie. // Get gets a cookie's value from its key. Use HasValueOrKey to check if a key or single-valued cookie is present in the cookie.
func (c *Cookie) Get(key string) []byte { func (c *Cookie) Get(key string) []byte { return c.kv.Get(key) }
nc := len(c.kvs)
for i := range nc {
kv := c.kvs[i]
if b2s(tok2bytes(c.buf, kv.key)) == key {
return tok2bytes(c.buf, kv.value)
}
}
return nil
}
// HasKeyOrSingleValue returns true if the cookie contains a pair with the given // HasKeyOrSingleValue returns true if the cookie contains a pair with the given
// key or a valueless attribute with the given text, i.e: "Secure" or "HttpOnly". // key or a valueless attribute with the given text, i.e: "Secure" or "HttpOnly".
func (c *Cookie) HasKeyOrSingleValue(keyOrSingleValue string) bool { func (c *Cookie) HasKeyOrSingleValue(keyOrSingleValue string) bool {
nc := len(c.kvs) return c.kv.Present(keyOrSingleValue)
for i := range nc {
kv := c.kvs[i]
if kv.key.len == 0 && b2s(tok2bytes(c.buf, kv.value)) == keyOrSingleValue ||
b2s(tok2bytes(c.buf, kv.key)) == keyOrSingleValue {
return true
}
}
return false
} }
// parseCookie parses a cookie inside cookie buffer and adds it to cookie buffer.. // parseCookie parses a cookie inside cookie buffer and adds it to cookie buffer..
@@ -162,16 +127,14 @@ func (c *Cookie) String() string {
// AppendKeyValues appends the HTTP header value of the cookie expected after the "Cookie:" string. Does not include trailing \r\n's. // AppendKeyValues appends the HTTP header value of the cookie expected after the "Cookie:" string. Does not include trailing \r\n's.
func (c *Cookie) AppendKeyValues(dst []byte) []byte { func (c *Cookie) AppendKeyValues(dst []byte) []byte {
nc := len(c.kvs) nc := c.kv.Len()
for i := range nc { for i := range nc {
kv := c.kvs[i] k, v := c.kv.At(i)
key := tok2bytes(c.buf, kv.key) if len(k) != 0 {
value := tok2bytes(c.buf, kv.value) dst = append(dst, k...)
if len(key) != 0 {
dst = append(dst, key...)
dst = append(dst, '=') dst = append(dst, '=')
} }
dst = append(dst, value...) dst = append(dst, v...)
if i+1 < nc { if i+1 < nc {
dst = append(dst, ';', ' ') dst = append(dst, ';', ' ')
} }
+1 -1
View File
@@ -62,7 +62,7 @@ func TestHeaderParseRequest(t *testing.T) {
} }
var c Cookie var c Cookie
cookie := hdr.Get("Cookie") cookie := hdr.Get("Cookie")
c.Reset(cookie) c.Reset(cookie, 0)
err = c.Parse() err = c.Parse()
if err != nil { if err != nil {
t.Error(err) t.Error(err)
+15 -1
View File
@@ -49,6 +49,18 @@ func (mb *KVBuffer) Get(key string) []byte {
return mb.musttoken(mb.kvs[v].value) return mb.musttoken(mb.kvs[v].value)
} }
// ForEach iterates over the cookie's key-value pairs until cb returns false.
func (c *KVBuffer) ForEach(cb func(key, value []byte) bool) {
nc := len(c.kvs)
for i := range nc {
if !c.kvs[i].isValid() {
continue
} else if !cb(c.At(i)) {
break
}
}
}
func (mb *KVBuffer) Present(key string) bool { func (mb *KVBuffer) Present(key string) bool {
return mb.getIdx(key) >= 0 return mb.getIdx(key) >= 0
} }
@@ -69,10 +81,12 @@ func (mb *KVBuffer) setInternal(key, value []byte) bool {
} }
func (mb *KVBuffer) Len() int { return len(mb.kvs) } func (mb *KVBuffer) Len() int { return len(mb.kvs) }
func (mb *KVBuffer) Pair(i int) (key, value []byte) { func (mb *KVBuffer) At(i int) (key, value []byte) {
kv := mb.kvs[i] kv := mb.kvs[i]
return mb.musttoken(kv.key), mb.musttoken(kv.value) return mb.musttoken(kv.key), mb.musttoken(kv.value)
} }
func (mb *KVBuffer) AtKey(i int) (key []byte) { return mb.musttoken(mb.kvs[i].key) }
func (mb *KVBuffer) AtValue(i int) (key []byte) { return mb.musttoken(mb.kvs[i].value) }
func (mb *KVBuffer) getIdx(key string) int { func (mb *KVBuffer) getIdx(key string) int {
for i, kv := range mb.kvs { for i, kv := range mb.kvs {
+2 -5
View File
@@ -356,13 +356,10 @@ func TestCookie_ForEach(t *testing.T) {
c.ParseBytes([]byte("a=1; b=2; c=3")) c.ParseBytes([]byte("a=1; b=2; c=3"))
var keys []string var keys []string
err := c.ForEach(func(key, value []byte) error { c.ForEach(func(key, value []byte) bool {
keys = append(keys, string(key)) keys = append(keys, string(key))
return nil return true
}) })
if err != nil {
t.Fatal(err)
}
if len(keys) != 3 { if len(keys) != 3 {
t.Fatalf("expected 3 cookie entries, got %d", len(keys)) t.Fatalf("expected 3 cookie entries, got %d", len(keys))
} }