well-goknown/vendor/github.com/valyala/fasthttp/args.go
2024-09-25 21:59:44 -04:00

644 lines
14 KiB
Go

package fasthttp
import (
"bytes"
"errors"
"io"
"sort"
"sync"
"github.com/valyala/bytebufferpool"
)
const (
argsNoValue = true
argsHasValue = false
)
// AcquireArgs returns an empty Args object from the pool.
//
// The returned Args may be returned to the pool with ReleaseArgs
// when no longer needed. This allows reducing GC load.
func AcquireArgs() *Args {
return argsPool.Get().(*Args)
}
// ReleaseArgs returns the object acquired via AcquireArgs to the pool.
//
// Do not access the released Args object, otherwise data races may occur.
func ReleaseArgs(a *Args) {
a.Reset()
argsPool.Put(a)
}
var argsPool = &sync.Pool{
New: func() any {
return &Args{}
},
}
// Args represents query arguments.
//
// It is forbidden copying Args instances. Create new instances instead
// and use CopyTo().
//
// Args instance MUST NOT be used from concurrently running goroutines.
type Args struct {
noCopy noCopy
args []argsKV
buf []byte
}
type argsKV struct {
key []byte
value []byte
noValue bool
}
// Reset clears query args.
func (a *Args) Reset() {
a.args = a.args[:0]
}
// CopyTo copies all args to dst.
func (a *Args) CopyTo(dst *Args) {
dst.args = copyArgs(dst.args, a.args)
}
// VisitAll calls f for each existing arg.
//
// f must not retain references to key and value after returning.
// Make key and/or value copies if you need storing them after returning.
func (a *Args) VisitAll(f func(key, value []byte)) {
visitArgs(a.args, f)
}
// Len returns the number of query args.
func (a *Args) Len() int {
return len(a.args)
}
// Parse parses the given string containing query args.
func (a *Args) Parse(s string) {
a.buf = append(a.buf[:0], s...)
a.ParseBytes(a.buf)
}
// ParseBytes parses the given b containing query args.
func (a *Args) ParseBytes(b []byte) {
a.Reset()
var s argsScanner
s.b = b
var kv *argsKV
a.args, kv = allocArg(a.args)
for s.next(kv) {
if len(kv.key) > 0 || len(kv.value) > 0 {
a.args, kv = allocArg(a.args)
}
}
a.args = releaseArg(a.args)
}
// String returns string representation of query args.
func (a *Args) String() string {
return string(a.QueryString())
}
// QueryString returns query string for the args.
//
// The returned value is valid until the Args is reused or released (ReleaseArgs).
// Do not store references to the returned value. Make copies instead.
func (a *Args) QueryString() []byte {
a.buf = a.AppendBytes(a.buf[:0])
return a.buf
}
// Sort sorts Args by key and then value using 'f' as comparison function.
//
// For example args.Sort(bytes.Compare).
func (a *Args) Sort(f func(x, y []byte) int) {
sort.SliceStable(a.args, func(i, j int) bool {
n := f(a.args[i].key, a.args[j].key)
if n == 0 {
return f(a.args[i].value, a.args[j].value) == -1
}
return n == -1
})
}
// AppendBytes appends query string to dst and returns the extended dst.
func (a *Args) AppendBytes(dst []byte) []byte {
for i, n := 0, len(a.args); i < n; i++ {
kv := &a.args[i]
dst = AppendQuotedArg(dst, kv.key)
if !kv.noValue {
dst = append(dst, '=')
if len(kv.value) > 0 {
dst = AppendQuotedArg(dst, kv.value)
}
}
if i+1 < n {
dst = append(dst, '&')
}
}
return dst
}
// WriteTo writes query string to w.
//
// WriteTo implements io.WriterTo interface.
func (a *Args) WriteTo(w io.Writer) (int64, error) {
n, err := w.Write(a.QueryString())
return int64(n), err
}
// Del deletes argument with the given key from query args.
func (a *Args) Del(key string) {
a.args = delAllArgs(a.args, key)
}
// DelBytes deletes argument with the given key from query args.
func (a *Args) DelBytes(key []byte) {
a.args = delAllArgs(a.args, b2s(key))
}
// Add adds 'key=value' argument.
//
// Multiple values for the same key may be added.
func (a *Args) Add(key, value string) {
a.args = appendArg(a.args, key, value, argsHasValue)
}
// AddBytesK adds 'key=value' argument.
//
// Multiple values for the same key may be added.
func (a *Args) AddBytesK(key []byte, value string) {
a.args = appendArg(a.args, b2s(key), value, argsHasValue)
}
// AddBytesV adds 'key=value' argument.
//
// Multiple values for the same key may be added.
func (a *Args) AddBytesV(key string, value []byte) {
a.args = appendArg(a.args, key, b2s(value), argsHasValue)
}
// AddBytesKV adds 'key=value' argument.
//
// Multiple values for the same key may be added.
func (a *Args) AddBytesKV(key, value []byte) {
a.args = appendArg(a.args, b2s(key), b2s(value), argsHasValue)
}
// AddNoValue adds only 'key' as argument without the '='.
//
// Multiple values for the same key may be added.
func (a *Args) AddNoValue(key string) {
a.args = appendArg(a.args, key, "", argsNoValue)
}
// AddBytesKNoValue adds only 'key' as argument without the '='.
//
// Multiple values for the same key may be added.
func (a *Args) AddBytesKNoValue(key []byte) {
a.args = appendArg(a.args, b2s(key), "", argsNoValue)
}
// Set sets 'key=value' argument.
func (a *Args) Set(key, value string) {
a.args = setArg(a.args, key, value, argsHasValue)
}
// SetBytesK sets 'key=value' argument.
func (a *Args) SetBytesK(key []byte, value string) {
a.args = setArg(a.args, b2s(key), value, argsHasValue)
}
// SetBytesV sets 'key=value' argument.
func (a *Args) SetBytesV(key string, value []byte) {
a.args = setArg(a.args, key, b2s(value), argsHasValue)
}
// SetBytesKV sets 'key=value' argument.
func (a *Args) SetBytesKV(key, value []byte) {
a.args = setArgBytes(a.args, key, value, argsHasValue)
}
// SetNoValue sets only 'key' as argument without the '='.
//
// Only key in argument, like key1&key2.
func (a *Args) SetNoValue(key string) {
a.args = setArg(a.args, key, "", argsNoValue)
}
// SetBytesKNoValue sets 'key' argument.
func (a *Args) SetBytesKNoValue(key []byte) {
a.args = setArg(a.args, b2s(key), "", argsNoValue)
}
// Peek returns query arg value for the given key.
//
// The returned value is valid until the Args is reused or released (ReleaseArgs).
// Do not store references to the returned value. Make copies instead.
func (a *Args) Peek(key string) []byte {
return peekArgStr(a.args, key)
}
// PeekBytes returns query arg value for the given key.
//
// The returned value is valid until the Args is reused or released (ReleaseArgs).
// Do not store references to the returned value. Make copies instead.
func (a *Args) PeekBytes(key []byte) []byte {
return peekArgBytes(a.args, key)
}
// PeekMulti returns all the arg values for the given key.
func (a *Args) PeekMulti(key string) [][]byte {
var values [][]byte
a.VisitAll(func(k, v []byte) {
if string(k) == key {
values = append(values, v)
}
})
return values
}
// PeekMultiBytes returns all the arg values for the given key.
func (a *Args) PeekMultiBytes(key []byte) [][]byte {
return a.PeekMulti(b2s(key))
}
// Has returns true if the given key exists in Args.
func (a *Args) Has(key string) bool {
return hasArg(a.args, key)
}
// HasBytes returns true if the given key exists in Args.
func (a *Args) HasBytes(key []byte) bool {
return hasArg(a.args, b2s(key))
}
// ErrNoArgValue is returned when Args value with the given key is missing.
var ErrNoArgValue = errors.New("no Args value for the given key")
// GetUint returns uint value for the given key.
func (a *Args) GetUint(key string) (int, error) {
value := a.Peek(key)
if len(value) == 0 {
return -1, ErrNoArgValue
}
return ParseUint(value)
}
// SetUint sets uint value for the given key.
func (a *Args) SetUint(key string, value int) {
bb := bytebufferpool.Get()
bb.B = AppendUint(bb.B[:0], value)
a.SetBytesV(key, bb.B)
bytebufferpool.Put(bb)
}
// SetUintBytes sets uint value for the given key.
func (a *Args) SetUintBytes(key []byte, value int) {
a.SetUint(b2s(key), value)
}
// GetUintOrZero returns uint value for the given key.
//
// Zero (0) is returned on error.
func (a *Args) GetUintOrZero(key string) int {
n, err := a.GetUint(key)
if err != nil {
n = 0
}
return n
}
// GetUfloat returns ufloat value for the given key.
func (a *Args) GetUfloat(key string) (float64, error) {
value := a.Peek(key)
if len(value) == 0 {
return -1, ErrNoArgValue
}
return ParseUfloat(value)
}
// GetUfloatOrZero returns ufloat value for the given key.
//
// Zero (0) is returned on error.
func (a *Args) GetUfloatOrZero(key string) float64 {
f, err := a.GetUfloat(key)
if err != nil {
f = 0
}
return f
}
// GetBool returns boolean value for the given key.
//
// true is returned for "1", "t", "T", "true", "TRUE", "True", "y", "yes", "Y", "YES", "Yes",
// otherwise false is returned.
func (a *Args) GetBool(key string) bool {
switch string(a.Peek(key)) {
// Support the same true cases as strconv.ParseBool
// See: https://github.com/golang/go/blob/4e1b11e2c9bdb0ddea1141eed487be1a626ff5be/src/strconv/atob.go#L12
// and Y and Yes versions.
case "1", "t", "T", "true", "TRUE", "True", "y", "yes", "Y", "YES", "Yes":
return true
default:
return false
}
}
func visitArgs(args []argsKV, f func(k, v []byte)) {
for i, n := 0, len(args); i < n; i++ {
kv := &args[i]
f(kv.key, kv.value)
}
}
func visitArgsKey(args []argsKV, f func(k []byte)) {
for i, n := 0, len(args); i < n; i++ {
kv := &args[i]
f(kv.key)
}
}
func copyArgs(dst, src []argsKV) []argsKV {
if cap(dst) < len(src) {
tmp := make([]argsKV, len(src))
dstLen := len(dst)
dst = dst[:cap(dst)] // copy all of dst.
copy(tmp, dst)
for i := dstLen; i < len(tmp); i++ {
// Make sure nothing is nil.
tmp[i].key = []byte{}
tmp[i].value = []byte{}
}
dst = tmp
}
n := len(src)
dst = dst[:n]
for i := 0; i < n; i++ {
dstKV := &dst[i]
srcKV := &src[i]
dstKV.key = append(dstKV.key[:0], srcKV.key...)
if srcKV.noValue {
dstKV.value = dstKV.value[:0]
} else {
dstKV.value = append(dstKV.value[:0], srcKV.value...)
}
dstKV.noValue = srcKV.noValue
}
return dst
}
func delAllArgsBytes(args []argsKV, key []byte) []argsKV {
return delAllArgs(args, b2s(key))
}
func delAllArgs(args []argsKV, key string) []argsKV {
for i, n := 0, len(args); i < n; i++ {
kv := &args[i]
if key == string(kv.key) {
tmp := *kv
copy(args[i:], args[i+1:])
n--
i--
args[n] = tmp
args = args[:n]
}
}
return args
}
func setArgBytes(h []argsKV, key, value []byte, noValue bool) []argsKV {
return setArg(h, b2s(key), b2s(value), noValue)
}
func setArg(h []argsKV, key, value string, noValue bool) []argsKV {
n := len(h)
for i := 0; i < n; i++ {
kv := &h[i]
if key == string(kv.key) {
if noValue {
kv.value = kv.value[:0]
} else {
kv.value = append(kv.value[:0], value...)
}
kv.noValue = noValue
return h
}
}
return appendArg(h, key, value, noValue)
}
func appendArgBytes(h []argsKV, key, value []byte, noValue bool) []argsKV {
return appendArg(h, b2s(key), b2s(value), noValue)
}
func appendArg(args []argsKV, key, value string, noValue bool) []argsKV {
var kv *argsKV
args, kv = allocArg(args)
kv.key = append(kv.key[:0], key...)
if noValue {
kv.value = kv.value[:0]
} else {
kv.value = append(kv.value[:0], value...)
}
kv.noValue = noValue
return args
}
func allocArg(h []argsKV) ([]argsKV, *argsKV) {
n := len(h)
if cap(h) > n {
h = h[:n+1]
} else {
h = append(h, argsKV{
value: []byte{},
})
}
return h, &h[n]
}
func releaseArg(h []argsKV) []argsKV {
return h[:len(h)-1]
}
func hasArg(h []argsKV, key string) bool {
for i, n := 0, len(h); i < n; i++ {
kv := &h[i]
if key == string(kv.key) {
return true
}
}
return false
}
func peekArgBytes(h []argsKV, k []byte) []byte {
for i, n := 0, len(h); i < n; i++ {
kv := &h[i]
if bytes.Equal(kv.key, k) {
return kv.value
}
}
return nil
}
func peekArgStr(h []argsKV, k string) []byte {
for i, n := 0, len(h); i < n; i++ {
kv := &h[i]
if string(kv.key) == k {
return kv.value
}
}
return nil
}
type argsScanner struct {
b []byte
}
func (s *argsScanner) next(kv *argsKV) bool {
if len(s.b) == 0 {
return false
}
kv.noValue = argsHasValue
isKey := true
k := 0
for i, c := range s.b {
switch c {
case '=':
if isKey {
isKey = false
kv.key = decodeArgAppend(kv.key[:0], s.b[:i])
k = i + 1
}
case '&':
if isKey {
kv.key = decodeArgAppend(kv.key[:0], s.b[:i])
kv.value = kv.value[:0]
kv.noValue = argsNoValue
} else {
kv.value = decodeArgAppend(kv.value[:0], s.b[k:i])
}
s.b = s.b[i+1:]
return true
}
}
if isKey {
kv.key = decodeArgAppend(kv.key[:0], s.b)
kv.value = kv.value[:0]
kv.noValue = argsNoValue
} else {
kv.value = decodeArgAppend(kv.value[:0], s.b[k:])
}
s.b = s.b[len(s.b):]
return true
}
func decodeArgAppend(dst, src []byte) []byte {
idxPercent := bytes.IndexByte(src, '%')
idxPlus := bytes.IndexByte(src, '+')
if idxPercent == -1 && idxPlus == -1 {
// fast path: src doesn't contain encoded chars
return append(dst, src...)
}
var idx int
switch {
case idxPercent == -1:
idx = idxPlus
case idxPlus == -1:
idx = idxPercent
case idxPercent > idxPlus:
idx = idxPlus
default:
idx = idxPercent
}
dst = append(dst, src[:idx]...)
// slow path
for i := idx; i < len(src); i++ {
c := src[i]
switch c {
case '%':
if i+2 >= len(src) {
return append(dst, src[i:]...)
}
x2 := hex2intTable[src[i+2]]
x1 := hex2intTable[src[i+1]]
if x1 == 16 || x2 == 16 {
dst = append(dst, '%')
} else {
dst = append(dst, x1<<4|x2)
i += 2
}
case '+':
dst = append(dst, ' ')
default:
dst = append(dst, c)
}
}
return dst
}
// decodeArgAppendNoPlus is almost identical to decodeArgAppend, but it doesn't
// substitute '+' with ' '.
//
// The function is copy-pasted from decodeArgAppend due to the performance
// reasons only.
func decodeArgAppendNoPlus(dst, src []byte) []byte {
idx := bytes.IndexByte(src, '%')
if idx < 0 {
// fast path: src doesn't contain encoded chars
return append(dst, src...)
}
dst = append(dst, src[:idx]...)
// slow path
for i := idx; i < len(src); i++ {
c := src[i]
if c == '%' {
if i+2 >= len(src) {
return append(dst, src[i:]...)
}
x2 := hex2intTable[src[i+2]]
x1 := hex2intTable[src[i+1]]
if x1 == 16 || x2 == 16 {
dst = append(dst, '%')
} else {
dst = append(dst, x1<<4|x2)
i += 2
}
} else {
dst = append(dst, c)
}
}
return dst
}
func peekAllArgBytesToDst(dst [][]byte, h []argsKV, k []byte) [][]byte {
for i, n := 0, len(h); i < n; i++ {
kv := &h[i]
if bytes.Equal(kv.key, k) {
dst = append(dst, kv.value)
}
}
return dst
}
func peekArgsKeys(dst [][]byte, h []argsKV) [][]byte {
for i, n := 0, len(h); i < n; i++ {
kv := &h[i]
dst = append(dst, kv.key)
}
return dst
}