args.go 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658
  1. package fasthttp
  2. import (
  3. "bytes"
  4. "errors"
  5. "io"
  6. "iter"
  7. "sort"
  8. "sync"
  9. )
  10. const (
  11. argsNoValue = true
  12. argsHasValue = false
  13. )
  14. // AcquireArgs returns an empty Args object from the pool.
  15. //
  16. // The returned Args may be returned to the pool with ReleaseArgs
  17. // when no longer needed. This allows reducing GC load.
  18. func AcquireArgs() *Args {
  19. return argsPool.Get().(*Args)
  20. }
  21. // ReleaseArgs returns the object acquired via AcquireArgs to the pool.
  22. //
  23. // Do not access the released Args object, otherwise data races may occur.
  24. func ReleaseArgs(a *Args) {
  25. a.Reset()
  26. argsPool.Put(a)
  27. }
  28. var argsPool = &sync.Pool{
  29. New: func() any {
  30. return &Args{}
  31. },
  32. }
  33. // Args represents query arguments.
  34. //
  35. // It is forbidden copying Args instances. Create new instances instead
  36. // and use CopyTo().
  37. //
  38. // Args instance MUST NOT be used from concurrently running goroutines.
  39. type Args struct {
  40. noCopy noCopy
  41. args []argsKV
  42. buf []byte
  43. }
  44. type argsKV struct {
  45. key []byte
  46. value []byte
  47. noValue bool
  48. }
  49. // Reset clears query args.
  50. func (a *Args) Reset() {
  51. a.args = a.args[:0]
  52. }
  53. // CopyTo copies all args to dst.
  54. func (a *Args) CopyTo(dst *Args) {
  55. dst.args = copyArgs(dst.args, a.args)
  56. }
  57. // All returns an iterator over key-value pairs from args.
  58. //
  59. // The key and value may invalid outside the iteration loop.
  60. // Make copies if you need to use them after the loop ends.
  61. //
  62. // Making modifications to the Args during the iteration loop leads to undefined
  63. // behavior and can cause panics.
  64. func (a *Args) All() iter.Seq2[[]byte, []byte] {
  65. return func(yield func([]byte, []byte) bool) {
  66. for i := range a.args {
  67. if !yield(a.args[i].key, a.args[i].value) {
  68. break
  69. }
  70. }
  71. }
  72. }
  73. // VisitAll calls f for each existing arg.
  74. //
  75. // f must not retain references to key and value after returning.
  76. // Make key and/or value copies if you need storing them after returning.
  77. //
  78. // Deprecated: Use All instead.
  79. func (a *Args) VisitAll(f func(key, value []byte)) {
  80. a.All()(func(key, value []byte) bool {
  81. f(key, value)
  82. return true
  83. })
  84. }
  85. // Len returns the number of query args.
  86. func (a *Args) Len() int {
  87. return len(a.args)
  88. }
  89. // Parse parses the given string containing query args.
  90. func (a *Args) Parse(s string) {
  91. a.buf = append(a.buf[:0], s...)
  92. a.ParseBytes(a.buf)
  93. }
  94. // ParseBytes parses the given b containing query args.
  95. func (a *Args) ParseBytes(b []byte) {
  96. a.Reset()
  97. var s argsScanner
  98. s.b = b
  99. var kv *argsKV
  100. a.args, kv = allocArg(a.args)
  101. for s.next(kv) {
  102. if len(kv.key) > 0 || len(kv.value) > 0 {
  103. a.args, kv = allocArg(a.args)
  104. }
  105. }
  106. a.args = releaseArg(a.args)
  107. }
  108. // String returns string representation of query args.
  109. func (a *Args) String() string {
  110. return string(a.QueryString())
  111. }
  112. // QueryString returns query string for the args.
  113. //
  114. // The returned value is valid until the Args is reused or released (ReleaseArgs).
  115. // Do not store references to the returned value. Make copies instead.
  116. func (a *Args) QueryString() []byte {
  117. a.buf = a.AppendBytes(a.buf[:0])
  118. return a.buf
  119. }
  120. // Sort sorts Args by key and then value using 'f' as comparison function.
  121. //
  122. // For example args.Sort(bytes.Compare).
  123. func (a *Args) Sort(f func(x, y []byte) int) {
  124. sort.SliceStable(a.args, func(i, j int) bool {
  125. n := f(a.args[i].key, a.args[j].key)
  126. if n == 0 {
  127. return f(a.args[i].value, a.args[j].value) == -1
  128. }
  129. return n == -1
  130. })
  131. }
  132. // SortKeys sorts Args by key only using 'f' as comparison function.
  133. //
  134. // For example args.SortKeys(bytes.Compare).
  135. func (a *Args) SortKeys(f func(x, y []byte) int) {
  136. sort.SliceStable(a.args, func(i, j int) bool {
  137. return f(a.args[i].key, a.args[j].key) == -1
  138. })
  139. }
  140. // AppendBytes appends query string to dst and returns the extended dst.
  141. func (a *Args) AppendBytes(dst []byte) []byte {
  142. for i, n := 0, len(a.args); i < n; i++ {
  143. kv := &a.args[i]
  144. dst = AppendQuotedArg(dst, kv.key)
  145. if !kv.noValue {
  146. dst = append(dst, '=')
  147. if len(kv.value) > 0 {
  148. dst = AppendQuotedArg(dst, kv.value)
  149. }
  150. }
  151. if i+1 < n {
  152. dst = append(dst, '&')
  153. }
  154. }
  155. return dst
  156. }
  157. // WriteTo writes query string to w.
  158. //
  159. // WriteTo implements io.WriterTo interface.
  160. func (a *Args) WriteTo(w io.Writer) (int64, error) {
  161. n, err := w.Write(a.QueryString())
  162. return int64(n), err
  163. }
  164. // Del deletes argument with the given key from query args.
  165. func (a *Args) Del(key string) {
  166. a.args = delAllArgsStable(a.args, key)
  167. }
  168. // DelBytes deletes argument with the given key from query args.
  169. func (a *Args) DelBytes(key []byte) {
  170. a.args = delAllArgsStable(a.args, b2s(key))
  171. }
  172. // Add adds 'key=value' argument.
  173. //
  174. // Multiple values for the same key may be added.
  175. func (a *Args) Add(key, value string) {
  176. a.args = appendArg(a.args, key, value, argsHasValue)
  177. }
  178. // AddBytesK adds 'key=value' argument.
  179. //
  180. // Multiple values for the same key may be added.
  181. func (a *Args) AddBytesK(key []byte, value string) {
  182. a.args = appendArg(a.args, b2s(key), value, argsHasValue)
  183. }
  184. // AddBytesV adds 'key=value' argument.
  185. //
  186. // Multiple values for the same key may be added.
  187. func (a *Args) AddBytesV(key string, value []byte) {
  188. a.args = appendArg(a.args, key, b2s(value), argsHasValue)
  189. }
  190. // AddBytesKV adds 'key=value' argument.
  191. //
  192. // Multiple values for the same key may be added.
  193. func (a *Args) AddBytesKV(key, value []byte) {
  194. a.args = appendArg(a.args, b2s(key), b2s(value), argsHasValue)
  195. }
  196. // AddNoValue adds only 'key' as argument without the '='.
  197. //
  198. // Multiple values for the same key may be added.
  199. func (a *Args) AddNoValue(key string) {
  200. a.args = appendArg(a.args, key, "", argsNoValue)
  201. }
  202. // AddBytesKNoValue adds only 'key' as argument without the '='.
  203. //
  204. // Multiple values for the same key may be added.
  205. func (a *Args) AddBytesKNoValue(key []byte) {
  206. a.args = appendArg(a.args, b2s(key), "", argsNoValue)
  207. }
  208. // Set sets 'key=value' argument.
  209. func (a *Args) Set(key, value string) {
  210. a.args = setArg(a.args, key, value, argsHasValue)
  211. }
  212. // SetBytesK sets 'key=value' argument.
  213. func (a *Args) SetBytesK(key []byte, value string) {
  214. a.args = setArg(a.args, b2s(key), value, argsHasValue)
  215. }
  216. // SetBytesV sets 'key=value' argument.
  217. func (a *Args) SetBytesV(key string, value []byte) {
  218. a.args = setArg(a.args, key, b2s(value), argsHasValue)
  219. }
  220. // SetBytesKV sets 'key=value' argument.
  221. func (a *Args) SetBytesKV(key, value []byte) {
  222. a.args = setArgBytes(a.args, key, value, argsHasValue)
  223. }
  224. // SetNoValue sets only 'key' as argument without the '='.
  225. //
  226. // Only key in argument, like key1&key2.
  227. func (a *Args) SetNoValue(key string) {
  228. a.args = setArg(a.args, key, "", argsNoValue)
  229. }
  230. // SetBytesKNoValue sets 'key' argument.
  231. func (a *Args) SetBytesKNoValue(key []byte) {
  232. a.args = setArg(a.args, b2s(key), "", argsNoValue)
  233. }
  234. // Peek returns query arg value for the given key.
  235. //
  236. // The returned value is valid until the Args is reused or released (ReleaseArgs).
  237. // Do not store references to the returned value. Make copies instead.
  238. func (a *Args) Peek(key string) []byte {
  239. return peekArgStr(a.args, key)
  240. }
  241. // PeekBytes returns query arg value for the given key.
  242. //
  243. // The returned value is valid until the Args is reused or released (ReleaseArgs).
  244. // Do not store references to the returned value. Make copies instead.
  245. func (a *Args) PeekBytes(key []byte) []byte {
  246. return peekArgBytes(a.args, key)
  247. }
  248. // PeekMulti returns all the arg values for the given key.
  249. func (a *Args) PeekMulti(key string) [][]byte {
  250. var values [][]byte
  251. for k, v := range a.All() {
  252. if string(k) == key {
  253. values = append(values, v)
  254. }
  255. }
  256. return values
  257. }
  258. // PeekMultiBytes returns all the arg values for the given key.
  259. func (a *Args) PeekMultiBytes(key []byte) [][]byte {
  260. return a.PeekMulti(b2s(key))
  261. }
  262. // Has returns true if the given key exists in Args.
  263. func (a *Args) Has(key string) bool {
  264. return hasArg(a.args, key)
  265. }
  266. // HasBytes returns true if the given key exists in Args.
  267. func (a *Args) HasBytes(key []byte) bool {
  268. return hasArg(a.args, b2s(key))
  269. }
  270. // ErrNoArgValue is returned when Args value with the given key is missing.
  271. var ErrNoArgValue = errors.New("no Args value for the given key")
  272. // GetUint returns uint value for the given key.
  273. func (a *Args) GetUint(key string) (int, error) {
  274. value := a.Peek(key)
  275. if len(value) == 0 {
  276. return -1, ErrNoArgValue
  277. }
  278. return ParseUint(value)
  279. }
  280. // SetUint sets uint value for the given key.
  281. func (a *Args) SetUint(key string, value int) {
  282. a.buf = AppendUint(a.buf[:0], value)
  283. a.SetBytesV(key, a.buf)
  284. }
  285. // SetUintBytes sets uint value for the given key.
  286. func (a *Args) SetUintBytes(key []byte, value int) {
  287. a.SetUint(b2s(key), value)
  288. }
  289. // GetUintOrZero returns uint value for the given key.
  290. //
  291. // Zero (0) is returned on error.
  292. func (a *Args) GetUintOrZero(key string) int {
  293. n, err := a.GetUint(key)
  294. if err != nil {
  295. n = 0
  296. }
  297. return n
  298. }
  299. // GetUfloat returns ufloat value for the given key.
  300. func (a *Args) GetUfloat(key string) (float64, error) {
  301. value := a.Peek(key)
  302. if len(value) == 0 {
  303. return -1, ErrNoArgValue
  304. }
  305. return ParseUfloat(value)
  306. }
  307. // GetUfloatOrZero returns ufloat value for the given key.
  308. //
  309. // Zero (0) is returned on error.
  310. func (a *Args) GetUfloatOrZero(key string) float64 {
  311. f, err := a.GetUfloat(key)
  312. if err != nil {
  313. f = 0
  314. }
  315. return f
  316. }
  317. // GetBool returns boolean value for the given key.
  318. //
  319. // true is returned for "1", "t", "T", "true", "TRUE", "True", "y", "yes", "Y", "YES", "Yes",
  320. // otherwise false is returned.
  321. func (a *Args) GetBool(key string) bool {
  322. switch string(a.Peek(key)) {
  323. // Support the same true cases as strconv.ParseBool
  324. // See: https://github.com/golang/go/blob/4e1b11e2c9bdb0ddea1141eed487be1a626ff5be/src/strconv/atob.go#L12
  325. // and Y and Yes versions.
  326. case "1", "t", "T", "true", "TRUE", "True", "y", "yes", "Y", "YES", "Yes":
  327. return true
  328. default:
  329. return false
  330. }
  331. }
  332. func copyArgs(dst, src []argsKV) []argsKV {
  333. if cap(dst) < len(src) {
  334. tmp := make([]argsKV, len(src))
  335. dstLen := len(dst)
  336. dst = dst[:cap(dst)] // copy all of dst.
  337. copy(tmp, dst)
  338. for i := dstLen; i < len(tmp); i++ {
  339. // Make sure nothing is nil.
  340. tmp[i].key = []byte{}
  341. tmp[i].value = []byte{}
  342. }
  343. dst = tmp
  344. }
  345. n := len(src)
  346. dst = dst[:n]
  347. for i := 0; i < n; i++ {
  348. dstKV := &dst[i]
  349. srcKV := &src[i]
  350. dstKV.key = append(dstKV.key[:0], srcKV.key...)
  351. if srcKV.noValue {
  352. dstKV.value = dstKV.value[:0]
  353. } else {
  354. dstKV.value = append(dstKV.value[:0], srcKV.value...)
  355. }
  356. dstKV.noValue = srcKV.noValue
  357. }
  358. return dst
  359. }
  360. func delAllArgsStable(args []argsKV, key string) []argsKV {
  361. for i, n := 0, len(args); i < n; i++ {
  362. kv := &args[i]
  363. if key == string(kv.key) {
  364. tmp := *kv
  365. copy(args[i:], args[i+1:])
  366. n--
  367. i--
  368. args[n] = tmp
  369. args = args[:n]
  370. }
  371. }
  372. return args
  373. }
  374. func delAllArgs(args []argsKV, key string) []argsKV {
  375. n := len(args)
  376. for i := 0; i < n; i++ {
  377. if key == string(args[i].key) {
  378. args[i], args[n-1] = args[n-1], args[i]
  379. n--
  380. i--
  381. }
  382. }
  383. return args[:n]
  384. }
  385. func setArgBytes(h []argsKV, key, value []byte, noValue bool) []argsKV {
  386. return setArg(h, b2s(key), b2s(value), noValue)
  387. }
  388. func setArg(h []argsKV, key, value string, noValue bool) []argsKV {
  389. n := len(h)
  390. for i := 0; i < n; i++ {
  391. kv := &h[i]
  392. if key == string(kv.key) {
  393. if noValue {
  394. kv.value = kv.value[:0]
  395. } else {
  396. kv.value = append(kv.value[:0], value...)
  397. }
  398. kv.noValue = noValue
  399. return h
  400. }
  401. }
  402. return appendArg(h, key, value, noValue)
  403. }
  404. func appendArgBytes(h []argsKV, key, value []byte, noValue bool) []argsKV {
  405. return appendArg(h, b2s(key), b2s(value), noValue)
  406. }
  407. func appendArg(args []argsKV, key, value string, noValue bool) []argsKV {
  408. var kv *argsKV
  409. args, kv = allocArg(args)
  410. kv.key = append(kv.key[:0], key...)
  411. if noValue {
  412. kv.value = kv.value[:0]
  413. } else {
  414. kv.value = append(kv.value[:0], value...)
  415. }
  416. kv.noValue = noValue
  417. return args
  418. }
  419. func allocArg(h []argsKV) ([]argsKV, *argsKV) {
  420. n := len(h)
  421. if cap(h) > n {
  422. h = h[:n+1]
  423. } else {
  424. h = append(h, argsKV{
  425. value: []byte{},
  426. })
  427. }
  428. return h, &h[n]
  429. }
  430. func releaseArg(h []argsKV) []argsKV {
  431. return h[:len(h)-1]
  432. }
  433. func hasArg(h []argsKV, key string) bool {
  434. for i, n := 0, len(h); i < n; i++ {
  435. kv := &h[i]
  436. if key == string(kv.key) {
  437. return true
  438. }
  439. }
  440. return false
  441. }
  442. func peekArgBytes(h []argsKV, k []byte) []byte {
  443. for i, n := 0, len(h); i < n; i++ {
  444. kv := &h[i]
  445. if bytes.Equal(kv.key, k) {
  446. return kv.value
  447. }
  448. }
  449. return nil
  450. }
  451. func peekArgStr(h []argsKV, k string) []byte {
  452. for i, n := 0, len(h); i < n; i++ {
  453. kv := &h[i]
  454. if string(kv.key) == k {
  455. return kv.value
  456. }
  457. }
  458. return nil
  459. }
  460. type argsScanner struct {
  461. b []byte
  462. }
  463. func (s *argsScanner) next(kv *argsKV) bool {
  464. if len(s.b) == 0 {
  465. return false
  466. }
  467. kv.noValue = argsHasValue
  468. isKey := true
  469. k := 0
  470. for i, c := range s.b {
  471. switch c {
  472. case '=':
  473. if isKey {
  474. isKey = false
  475. kv.key = decodeArgAppend(kv.key[:0], s.b[:i])
  476. k = i + 1
  477. }
  478. case '&':
  479. if isKey {
  480. kv.key = decodeArgAppend(kv.key[:0], s.b[:i])
  481. kv.value = kv.value[:0]
  482. kv.noValue = argsNoValue
  483. } else {
  484. kv.value = decodeArgAppend(kv.value[:0], s.b[k:i])
  485. }
  486. s.b = s.b[i+1:]
  487. return true
  488. }
  489. }
  490. if isKey {
  491. kv.key = decodeArgAppend(kv.key[:0], s.b)
  492. kv.value = kv.value[:0]
  493. kv.noValue = argsNoValue
  494. } else {
  495. kv.value = decodeArgAppend(kv.value[:0], s.b[k:])
  496. }
  497. s.b = s.b[len(s.b):]
  498. return true
  499. }
  500. func decodeArgAppend(dst, src []byte) []byte {
  501. idxPercent := bytes.IndexByte(src, '%')
  502. idxPlus := bytes.IndexByte(src, '+')
  503. if idxPercent == -1 && idxPlus == -1 {
  504. // fast path: src doesn't contain encoded chars
  505. return append(dst, src...)
  506. }
  507. var idx int
  508. switch {
  509. case idxPercent == -1:
  510. idx = idxPlus
  511. case idxPlus == -1:
  512. idx = idxPercent
  513. case idxPercent > idxPlus:
  514. idx = idxPlus
  515. default:
  516. idx = idxPercent
  517. }
  518. dst = append(dst, src[:idx]...)
  519. // slow path
  520. for i := idx; i < len(src); i++ {
  521. c := src[i]
  522. switch c {
  523. case '%':
  524. if i+2 >= len(src) {
  525. return append(dst, src[i:]...)
  526. }
  527. x2 := hex2intTable[src[i+2]]
  528. x1 := hex2intTable[src[i+1]]
  529. if x1 == 16 || x2 == 16 {
  530. dst = append(dst, '%')
  531. } else {
  532. dst = append(dst, x1<<4|x2)
  533. i += 2
  534. }
  535. case '+':
  536. dst = append(dst, ' ')
  537. default:
  538. dst = append(dst, c)
  539. }
  540. }
  541. return dst
  542. }
  543. // decodeArgAppendNoPlus is almost identical to decodeArgAppend, but it doesn't
  544. // substitute '+' with ' '.
  545. //
  546. // The function is copy-pasted from decodeArgAppend due to the performance
  547. // reasons only.
  548. func decodeArgAppendNoPlus(dst, src []byte) []byte {
  549. idx := bytes.IndexByte(src, '%')
  550. if idx < 0 {
  551. // fast path: src doesn't contain encoded chars
  552. return append(dst, src...)
  553. }
  554. dst = append(dst, src[:idx]...)
  555. // slow path
  556. for i := idx; i < len(src); i++ {
  557. c := src[i]
  558. if c == '%' {
  559. if i+2 >= len(src) {
  560. return append(dst, src[i:]...)
  561. }
  562. x2 := hex2intTable[src[i+2]]
  563. x1 := hex2intTable[src[i+1]]
  564. if x1 == 16 || x2 == 16 {
  565. dst = append(dst, '%')
  566. } else {
  567. dst = append(dst, x1<<4|x2)
  568. i += 2
  569. }
  570. } else {
  571. dst = append(dst, c)
  572. }
  573. }
  574. return dst
  575. }
  576. func peekAllArgBytesToDst(dst [][]byte, h []argsKV, k []byte) [][]byte {
  577. for i, n := 0, len(h); i < n; i++ {
  578. kv := &h[i]
  579. if bytes.Equal(kv.key, k) {
  580. dst = append(dst, kv.value)
  581. }
  582. }
  583. return dst
  584. }