retriablerequest.go 1.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. package autorest
  2. // Copyright 2017 Microsoft Corporation
  3. //
  4. // Licensed under the Apache License, Version 2.0 (the "License");
  5. // you may not use this file except in compliance with the License.
  6. // You may obtain a copy of the License at
  7. //
  8. // http://www.apache.org/licenses/LICENSE-2.0
  9. //
  10. // Unless required by applicable law or agreed to in writing, software
  11. // distributed under the License is distributed on an "AS IS" BASIS,
  12. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. // See the License for the specific language governing permissions and
  14. // limitations under the License.
  15. import (
  16. "bytes"
  17. "io"
  18. "io/ioutil"
  19. "net/http"
  20. )
  21. // NewRetriableRequest returns a wrapper around an HTTP request that support retry logic.
  22. func NewRetriableRequest(req *http.Request) *RetriableRequest {
  23. return &RetriableRequest{req: req}
  24. }
  25. // Request returns the wrapped HTTP request.
  26. func (rr *RetriableRequest) Request() *http.Request {
  27. return rr.req
  28. }
  29. func (rr *RetriableRequest) prepareFromByteReader() (err error) {
  30. // fall back to making a copy (only do this once)
  31. b := []byte{}
  32. if rr.req.ContentLength > 0 {
  33. b = make([]byte, rr.req.ContentLength)
  34. _, err = io.ReadFull(rr.req.Body, b)
  35. if err != nil {
  36. return err
  37. }
  38. } else {
  39. b, err = ioutil.ReadAll(rr.req.Body)
  40. if err != nil {
  41. return err
  42. }
  43. }
  44. rr.br = bytes.NewReader(b)
  45. rr.req.Body = ioutil.NopCloser(rr.br)
  46. return err
  47. }