You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

115 lines
2.8 KiB

  1. package util
  2. import (
  3. "bytes"
  4. "fmt"
  5. "github.com/valyala/fasthttp"
  6. "sync"
  7. "time"
  8. )
  9. var (
  10. fastClient = &fasthttp.Client{
  11. NoDefaultUserAgentHeader: true, // Don't send: User-Agent: fasthttp
  12. MaxConnsPerHost: 1024,
  13. ReadBufferSize: 4096, // Make sure to set this big enough that your whole request can be read at once.
  14. WriteBufferSize: 64 * 1024, // Same but for your response.
  15. ReadTimeout: time.Second,
  16. WriteTimeout: time.Second,
  17. MaxIdleConnDuration: time.Minute,
  18. DisableHeaderNamesNormalizing: true, // If you set the case on your headers correctly you can enable this.
  19. }
  20. // Put everything in pools to prevent garbage.
  21. bytesPool = sync.Pool{
  22. New: func() interface{} {
  23. b := make([]byte, 0)
  24. return &b
  25. },
  26. }
  27. responsePool = sync.Pool{
  28. New: func() interface{} {
  29. return make(chan *fasthttp.Response)
  30. },
  31. }
  32. )
  33. func FastGet(url string) ([]byte, bool, error) {
  34. req := fasthttp.AcquireRequest()
  35. res := fasthttp.AcquireResponse()
  36. defer fasthttp.ReleaseRequest(req)
  37. defer fasthttp.ReleaseResponse(res)
  38. req.SetRequestURIBytes([]byte(url))
  39. req.Header.Add("Accept-Encoding", "gzip")
  40. err := fastClient.Do(req, res)
  41. if err != nil {
  42. return nil, true, err
  43. }
  44. var data []byte
  45. contentEncoding := res.Header.Peek("Content-Encoding")
  46. if bytes.Compare(contentEncoding, []byte("gzip")) == 0 {
  47. data, err = res.BodyGunzip()
  48. } else {
  49. data = res.Body()
  50. }
  51. out := make([]byte, len(data))
  52. copy(out, data)
  53. if res.StatusCode() >= 400 {
  54. retryable := res.StatusCode() >= 500
  55. return nil, retryable, fmt.Errorf("%s: %d", url, res.StatusCode())
  56. }
  57. if err != nil {
  58. return nil, false, err
  59. }
  60. return out, false, nil
  61. }
  62. func FastReadUrlAsStream(fileUrl string, cipherKey []byte, isContentGzipped bool, isFullChunk bool, offset int64, size int, fn func(data []byte)) (retryable bool, err error) {
  63. if cipherKey != nil {
  64. return readEncryptedUrl(fileUrl, cipherKey, isContentGzipped, isFullChunk, offset, size, fn)
  65. }
  66. req := fasthttp.AcquireRequest()
  67. res := fasthttp.AcquireResponse()
  68. defer fasthttp.ReleaseRequest(req)
  69. defer fasthttp.ReleaseResponse(res)
  70. req.SetRequestURIBytes([]byte(fileUrl))
  71. if isFullChunk {
  72. req.Header.Add("Accept-Encoding", "gzip")
  73. } else {
  74. req.Header.Add("Range", fmt.Sprintf("bytes=%d-%d", offset, offset+int64(size)-1))
  75. }
  76. if err = fastClient.Do(req, res); err != nil {
  77. return true, err
  78. }
  79. if res.StatusCode() >= 400 {
  80. retryable = res.StatusCode() >= 500
  81. return retryable, fmt.Errorf("%s: %d", fileUrl, res.StatusCode())
  82. }
  83. contentEncoding := res.Header.Peek("Content-Encoding")
  84. if bytes.Compare(contentEncoding, []byte("gzip")) == 0 {
  85. bodyData, err := res.BodyGunzip()
  86. if err != nil {
  87. return false, err
  88. }
  89. fn(bodyData)
  90. } else {
  91. fn(res.Body())
  92. }
  93. return false, nil
  94. }