args.go 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589
  1. package fasthttp
  2. import (
  3. "bytes"
  4. "errors"
  5. "io"
  6. "sort"
  7. "sync"
  8. "github.com/valyala/bytebufferpool"
  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() interface{} {
  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.Reset()
  56. dst.args = copyArgs(dst.args, a.args)
  57. }
  58. // VisitAll calls f for each existing arg.
  59. //
  60. // f must not retain references to key and value after returning.
  61. // Make key and/or value copies if you need storing them after returning.
  62. func (a *Args) VisitAll(f func(key, value []byte)) {
  63. visitArgs(a.args, f)
  64. }
  65. // Len returns the number of query args.
  66. func (a *Args) Len() int {
  67. return len(a.args)
  68. }
  69. // Parse parses the given string containing query args.
  70. func (a *Args) Parse(s string) {
  71. a.buf = append(a.buf[:0], s...)
  72. a.ParseBytes(a.buf)
  73. }
  74. // ParseBytes parses the given b containing query args.
  75. func (a *Args) ParseBytes(b []byte) {
  76. a.Reset()
  77. var s argsScanner
  78. s.b = b
  79. var kv *argsKV
  80. a.args, kv = allocArg(a.args)
  81. for s.next(kv) {
  82. if len(kv.key) > 0 || len(kv.value) > 0 {
  83. a.args, kv = allocArg(a.args)
  84. }
  85. }
  86. a.args = releaseArg(a.args)
  87. }
  88. // String returns string representation of query args.
  89. func (a *Args) String() string {
  90. return string(a.QueryString())
  91. }
  92. // QueryString returns query string for the args.
  93. //
  94. // The returned value is valid until the next call to Args methods.
  95. func (a *Args) QueryString() []byte {
  96. a.buf = a.AppendBytes(a.buf[:0])
  97. return a.buf
  98. }
  99. // Sort sorts Args by key and then value using 'f' as comparison function.
  100. //
  101. // For example args.Sort(bytes.Compare)
  102. func (a *Args) Sort(f func(x, y []byte) int) {
  103. sort.SliceStable(a.args, func(i, j int) bool {
  104. n := f(a.args[i].key, a.args[j].key)
  105. if n == 0 {
  106. return f(a.args[i].value, a.args[j].value) == -1
  107. }
  108. return n == -1
  109. })
  110. }
  111. // AppendBytes appends query string to dst and returns the extended dst.
  112. func (a *Args) AppendBytes(dst []byte) []byte {
  113. for i, n := 0, len(a.args); i < n; i++ {
  114. kv := &a.args[i]
  115. dst = AppendQuotedArg(dst, kv.key)
  116. if !kv.noValue {
  117. dst = append(dst, '=')
  118. if len(kv.value) > 0 {
  119. dst = AppendQuotedArg(dst, kv.value)
  120. }
  121. }
  122. if i+1 < n {
  123. dst = append(dst, '&')
  124. }
  125. }
  126. return dst
  127. }
  128. // WriteTo writes query string to w.
  129. //
  130. // WriteTo implements io.WriterTo interface.
  131. func (a *Args) WriteTo(w io.Writer) (int64, error) {
  132. n, err := w.Write(a.QueryString())
  133. return int64(n), err
  134. }
  135. // Del deletes argument with the given key from query args.
  136. func (a *Args) Del(key string) {
  137. a.args = delAllArgs(a.args, key)
  138. }
  139. // DelBytes deletes argument with the given key from query args.
  140. func (a *Args) DelBytes(key []byte) {
  141. a.args = delAllArgs(a.args, b2s(key))
  142. }
  143. // Add adds 'key=value' argument.
  144. //
  145. // Multiple values for the same key may be added.
  146. func (a *Args) Add(key, value string) {
  147. a.args = appendArg(a.args, key, value, argsHasValue)
  148. }
  149. // AddBytesK adds 'key=value' argument.
  150. //
  151. // Multiple values for the same key may be added.
  152. func (a *Args) AddBytesK(key []byte, value string) {
  153. a.args = appendArg(a.args, b2s(key), value, argsHasValue)
  154. }
  155. // AddBytesV adds 'key=value' argument.
  156. //
  157. // Multiple values for the same key may be added.
  158. func (a *Args) AddBytesV(key string, value []byte) {
  159. a.args = appendArg(a.args, key, b2s(value), argsHasValue)
  160. }
  161. // AddBytesKV adds 'key=value' argument.
  162. //
  163. // Multiple values for the same key may be added.
  164. func (a *Args) AddBytesKV(key, value []byte) {
  165. a.args = appendArg(a.args, b2s(key), b2s(value), argsHasValue)
  166. }
  167. // AddNoValue adds only 'key' as argument without the '='.
  168. //
  169. // Multiple values for the same key may be added.
  170. func (a *Args) AddNoValue(key string) {
  171. a.args = appendArg(a.args, key, "", argsNoValue)
  172. }
  173. // AddBytesKNoValue adds only 'key' as argument without the '='.
  174. //
  175. // Multiple values for the same key may be added.
  176. func (a *Args) AddBytesKNoValue(key []byte) {
  177. a.args = appendArg(a.args, b2s(key), "", argsNoValue)
  178. }
  179. // Set sets 'key=value' argument.
  180. func (a *Args) Set(key, value string) {
  181. a.args = setArg(a.args, key, value, argsHasValue)
  182. }
  183. // SetBytesK sets 'key=value' argument.
  184. func (a *Args) SetBytesK(key []byte, value string) {
  185. a.args = setArg(a.args, b2s(key), value, argsHasValue)
  186. }
  187. // SetBytesV sets 'key=value' argument.
  188. func (a *Args) SetBytesV(key string, value []byte) {
  189. a.args = setArg(a.args, key, b2s(value), argsHasValue)
  190. }
  191. // SetBytesKV sets 'key=value' argument.
  192. func (a *Args) SetBytesKV(key, value []byte) {
  193. a.args = setArgBytes(a.args, key, value, argsHasValue)
  194. }
  195. // SetNoValue sets only 'key' as argument without the '='.
  196. //
  197. // Only key in argumemt, like key1&key2
  198. func (a *Args) SetNoValue(key string) {
  199. a.args = setArg(a.args, key, "", argsNoValue)
  200. }
  201. // SetBytesKNoValue sets 'key' argument.
  202. func (a *Args) SetBytesKNoValue(key []byte) {
  203. a.args = setArg(a.args, b2s(key), "", argsNoValue)
  204. }
  205. // Peek returns query arg value for the given key.
  206. //
  207. // Returned value is valid until the next Args call.
  208. func (a *Args) Peek(key string) []byte {
  209. return peekArgStr(a.args, key)
  210. }
  211. // PeekBytes returns query arg value for the given key.
  212. //
  213. // Returned value is valid until the next Args call.
  214. func (a *Args) PeekBytes(key []byte) []byte {
  215. return peekArgBytes(a.args, key)
  216. }
  217. // PeekMulti returns all the arg values for the given key.
  218. func (a *Args) PeekMulti(key string) [][]byte {
  219. var values [][]byte
  220. a.VisitAll(func(k, v []byte) {
  221. if string(k) == key {
  222. values = append(values, v)
  223. }
  224. })
  225. return values
  226. }
  227. // PeekMultiBytes returns all the arg values for the given key.
  228. func (a *Args) PeekMultiBytes(key []byte) [][]byte {
  229. return a.PeekMulti(b2s(key))
  230. }
  231. // Has returns true if the given key exists in Args.
  232. func (a *Args) Has(key string) bool {
  233. return hasArg(a.args, key)
  234. }
  235. // HasBytes returns true if the given key exists in Args.
  236. func (a *Args) HasBytes(key []byte) bool {
  237. return hasArg(a.args, b2s(key))
  238. }
  239. // ErrNoArgValue is returned when Args value with the given key is missing.
  240. var ErrNoArgValue = errors.New("no Args value for the given key")
  241. // GetUint returns uint value for the given key.
  242. func (a *Args) GetUint(key string) (int, error) {
  243. value := a.Peek(key)
  244. if len(value) == 0 {
  245. return -1, ErrNoArgValue
  246. }
  247. return ParseUint(value)
  248. }
  249. // SetUint sets uint value for the given key.
  250. func (a *Args) SetUint(key string, value int) {
  251. bb := bytebufferpool.Get()
  252. bb.B = AppendUint(bb.B[:0], value)
  253. a.SetBytesV(key, bb.B)
  254. bytebufferpool.Put(bb)
  255. }
  256. // SetUintBytes sets uint value for the given key.
  257. func (a *Args) SetUintBytes(key []byte, value int) {
  258. a.SetUint(b2s(key), value)
  259. }
  260. // GetUintOrZero returns uint value for the given key.
  261. //
  262. // Zero (0) is returned on error.
  263. func (a *Args) GetUintOrZero(key string) int {
  264. n, err := a.GetUint(key)
  265. if err != nil {
  266. n = 0
  267. }
  268. return n
  269. }
  270. // GetUfloat returns ufloat value for the given key.
  271. func (a *Args) GetUfloat(key string) (float64, error) {
  272. value := a.Peek(key)
  273. if len(value) == 0 {
  274. return -1, ErrNoArgValue
  275. }
  276. return ParseUfloat(value)
  277. }
  278. // GetUfloatOrZero returns ufloat value for the given key.
  279. //
  280. // Zero (0) is returned on error.
  281. func (a *Args) GetUfloatOrZero(key string) float64 {
  282. f, err := a.GetUfloat(key)
  283. if err != nil {
  284. f = 0
  285. }
  286. return f
  287. }
  288. // GetBool returns boolean value for the given key.
  289. //
  290. // true is returned for "1", "t", "T", "true", "TRUE", "True", "y", "yes", "Y", "YES", "Yes",
  291. // otherwise false is returned.
  292. func (a *Args) GetBool(key string) bool {
  293. switch b2s(a.Peek(key)) {
  294. // Support the same true cases as strconv.ParseBool
  295. // See: https://github.com/golang/go/blob/4e1b11e2c9bdb0ddea1141eed487be1a626ff5be/src/strconv/atob.go#L12
  296. // and Y and Yes versions.
  297. case "1", "t", "T", "true", "TRUE", "True", "y", "yes", "Y", "YES", "Yes":
  298. return true
  299. default:
  300. return false
  301. }
  302. }
  303. func visitArgs(args []argsKV, f func(k, v []byte)) {
  304. for i, n := 0, len(args); i < n; i++ {
  305. kv := &args[i]
  306. f(kv.key, kv.value)
  307. }
  308. }
  309. func copyArgs(dst, src []argsKV) []argsKV {
  310. if cap(dst) < len(src) {
  311. tmp := make([]argsKV, len(src))
  312. copy(tmp, dst)
  313. dst = tmp
  314. }
  315. n := len(src)
  316. dst = dst[:n]
  317. for i := 0; i < n; i++ {
  318. dstKV := &dst[i]
  319. srcKV := &src[i]
  320. dstKV.key = append(dstKV.key[:0], srcKV.key...)
  321. if srcKV.noValue {
  322. dstKV.value = dstKV.value[:0]
  323. } else {
  324. dstKV.value = append(dstKV.value[:0], srcKV.value...)
  325. }
  326. dstKV.noValue = srcKV.noValue
  327. }
  328. return dst
  329. }
  330. func delAllArgsBytes(args []argsKV, key []byte) []argsKV {
  331. return delAllArgs(args, b2s(key))
  332. }
  333. func delAllArgs(args []argsKV, key string) []argsKV {
  334. for i, n := 0, len(args); i < n; i++ {
  335. kv := &args[i]
  336. if key == string(kv.key) {
  337. tmp := *kv
  338. copy(args[i:], args[i+1:])
  339. n--
  340. args[n] = tmp
  341. args = args[:n]
  342. }
  343. }
  344. return args
  345. }
  346. func setArgBytes(h []argsKV, key, value []byte, noValue bool) []argsKV {
  347. return setArg(h, b2s(key), b2s(value), noValue)
  348. }
  349. func setArg(h []argsKV, key, value string, noValue bool) []argsKV {
  350. n := len(h)
  351. for i := 0; i < n; i++ {
  352. kv := &h[i]
  353. if key == string(kv.key) {
  354. if noValue {
  355. kv.value = kv.value[:0]
  356. } else {
  357. kv.value = append(kv.value[:0], value...)
  358. }
  359. kv.noValue = noValue
  360. return h
  361. }
  362. }
  363. return appendArg(h, key, value, noValue)
  364. }
  365. func appendArgBytes(h []argsKV, key, value []byte, noValue bool) []argsKV {
  366. return appendArg(h, b2s(key), b2s(value), noValue)
  367. }
  368. func appendArg(args []argsKV, key, value string, noValue bool) []argsKV {
  369. var kv *argsKV
  370. args, kv = allocArg(args)
  371. kv.key = append(kv.key[:0], key...)
  372. if noValue {
  373. kv.value = kv.value[:0]
  374. } else {
  375. kv.value = append(kv.value[:0], value...)
  376. }
  377. kv.noValue = noValue
  378. return args
  379. }
  380. func allocArg(h []argsKV) ([]argsKV, *argsKV) {
  381. n := len(h)
  382. if cap(h) > n {
  383. h = h[:n+1]
  384. } else {
  385. h = append(h, argsKV{})
  386. }
  387. return h, &h[n]
  388. }
  389. func releaseArg(h []argsKV) []argsKV {
  390. return h[:len(h)-1]
  391. }
  392. func hasArg(h []argsKV, key string) bool {
  393. for i, n := 0, len(h); i < n; i++ {
  394. kv := &h[i]
  395. if key == string(kv.key) {
  396. return true
  397. }
  398. }
  399. return false
  400. }
  401. func peekArgBytes(h []argsKV, k []byte) []byte {
  402. for i, n := 0, len(h); i < n; i++ {
  403. kv := &h[i]
  404. if bytes.Equal(kv.key, k) {
  405. return kv.value
  406. }
  407. }
  408. return nil
  409. }
  410. func peekArgStr(h []argsKV, k string) []byte {
  411. for i, n := 0, len(h); i < n; i++ {
  412. kv := &h[i]
  413. if string(kv.key) == k {
  414. return kv.value
  415. }
  416. }
  417. return nil
  418. }
  419. type argsScanner struct {
  420. b []byte
  421. }
  422. func (s *argsScanner) next(kv *argsKV) bool {
  423. if len(s.b) == 0 {
  424. return false
  425. }
  426. kv.noValue = argsHasValue
  427. isKey := true
  428. k := 0
  429. for i, c := range s.b {
  430. switch c {
  431. case '=':
  432. if isKey {
  433. isKey = false
  434. kv.key = decodeArgAppend(kv.key[:0], s.b[:i])
  435. k = i + 1
  436. }
  437. case '&':
  438. if isKey {
  439. kv.key = decodeArgAppend(kv.key[:0], s.b[:i])
  440. kv.value = kv.value[:0]
  441. kv.noValue = argsNoValue
  442. } else {
  443. kv.value = decodeArgAppend(kv.value[:0], s.b[k:i])
  444. }
  445. s.b = s.b[i+1:]
  446. return true
  447. }
  448. }
  449. if isKey {
  450. kv.key = decodeArgAppend(kv.key[:0], s.b)
  451. kv.value = kv.value[:0]
  452. kv.noValue = argsNoValue
  453. } else {
  454. kv.value = decodeArgAppend(kv.value[:0], s.b[k:])
  455. }
  456. s.b = s.b[len(s.b):]
  457. return true
  458. }
  459. func decodeArgAppend(dst, src []byte) []byte {
  460. if bytes.IndexByte(src, '%') < 0 && bytes.IndexByte(src, '+') < 0 {
  461. // fast path: src doesn't contain encoded chars
  462. return append(dst, src...)
  463. }
  464. // slow path
  465. for i := 0; i < len(src); i++ {
  466. c := src[i]
  467. if c == '%' {
  468. if i+2 >= len(src) {
  469. return append(dst, src[i:]...)
  470. }
  471. x2 := hex2intTable[src[i+2]]
  472. x1 := hex2intTable[src[i+1]]
  473. if x1 == 16 || x2 == 16 {
  474. dst = append(dst, '%')
  475. } else {
  476. dst = append(dst, x1<<4|x2)
  477. i += 2
  478. }
  479. } else if c == '+' {
  480. dst = append(dst, ' ')
  481. } else {
  482. dst = append(dst, c)
  483. }
  484. }
  485. return dst
  486. }
  487. // decodeArgAppendNoPlus is almost identical to decodeArgAppend, but it doesn't
  488. // substitute '+' with ' '.
  489. //
  490. // The function is copy-pasted from decodeArgAppend due to the performance
  491. // reasons only.
  492. func decodeArgAppendNoPlus(dst, src []byte) []byte {
  493. if bytes.IndexByte(src, '%') < 0 {
  494. // fast path: src doesn't contain encoded chars
  495. return append(dst, src...)
  496. }
  497. // slow path
  498. for i := 0; i < len(src); i++ {
  499. c := src[i]
  500. if c == '%' {
  501. if i+2 >= len(src) {
  502. return append(dst, src[i:]...)
  503. }
  504. x2 := hex2intTable[src[i+2]]
  505. x1 := hex2intTable[src[i+1]]
  506. if x1 == 16 || x2 == 16 {
  507. dst = append(dst, '%')
  508. } else {
  509. dst = append(dst, x1<<4|x2)
  510. i += 2
  511. }
  512. } else {
  513. dst = append(dst, c)
  514. }
  515. }
  516. return dst
  517. }