client.go 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553
  1. package client
  2. import (
  3. "bytes"
  4. "context"
  5. "encoding/json"
  6. "errors"
  7. "fmt"
  8. "io"
  9. "mime/multipart"
  10. "net/http"
  11. "net/url"
  12. "os"
  13. "strconv"
  14. "strings"
  15. "golang.org/x/time/rate"
  16. )
  17. // A Client is an HTTP client. Initialize with the New package-level function.
  18. type Client struct {
  19. opts []Option // keep for clones.
  20. HTTPClient *http.Client
  21. // BaseURL prepends to all requests.
  22. BaseURL string
  23. // A list of persistent request options.
  24. PersistentRequestOptions []RequestOption
  25. // Optional rate limiter instance initialized by the RateLimit method.
  26. rateLimiter *rate.Limiter
  27. // Optional handlers that are being fired before and after each new request.
  28. requestHandlers []RequestHandler
  29. // store it here for future use.
  30. keepAlive bool
  31. }
  32. // New returns a new Iris HTTP Client.
  33. // Available options:
  34. // - BaseURL
  35. // - Timeout
  36. // - PersistentRequestOptions
  37. // - RateLimit
  38. //
  39. // Look the Client.Do/JSON/... methods to send requests and
  40. // ReadXXX methods to read responses.
  41. //
  42. // The default content type to send and receive data is JSON.
  43. func New(opts ...Option) *Client {
  44. c := &Client{
  45. opts: opts,
  46. HTTPClient: &http.Client{},
  47. PersistentRequestOptions: defaultRequestOptions,
  48. requestHandlers: defaultRequestHandlers,
  49. }
  50. for _, opt := range c.opts { // c.opts in order to make with `NoOption` work.
  51. opt(c)
  52. }
  53. if transport, ok := c.HTTPClient.Transport.(*http.Transport); ok {
  54. c.keepAlive = !transport.DisableKeepAlives
  55. }
  56. return c
  57. }
  58. // NoOption is a helper function that clears the previous options in the chain.
  59. // See `Client.Clone` method.
  60. var NoOption = func(c *Client) { c.opts = make([]Option, 0) /* clear previous options */ }
  61. // Clone returns a new Client with the same options as the original.
  62. // If you want to override the options from the base "c" Client,
  63. // use the `NoOption` variable as the 1st argument.
  64. func (c *Client) Clone(opts ...Option) *Client {
  65. return New(append(c.opts, opts...)...)
  66. }
  67. // RegisterRequestHandler registers one or more request handlers
  68. // to be ran before and after of each new request.
  69. //
  70. // Request handler's BeginRequest method run after each request constructed
  71. // and right before sent to the server.
  72. //
  73. // Request handler's EndRequest method run after response each received
  74. // and right before methods return back to the caller.
  75. //
  76. // Any request handlers MUST be set right after the Client's initialization.
  77. func (c *Client) RegisterRequestHandler(reqHandlers ...RequestHandler) {
  78. reqHandlersToRegister := make([]RequestHandler, 0, len(reqHandlers))
  79. for _, h := range reqHandlers {
  80. if h == nil {
  81. continue
  82. }
  83. reqHandlersToRegister = append(reqHandlersToRegister, h)
  84. }
  85. c.requestHandlers = append(c.requestHandlers, reqHandlersToRegister...)
  86. }
  87. func (c *Client) emitBeginRequest(ctx context.Context, req *http.Request) error {
  88. if len(c.requestHandlers) == 0 {
  89. return nil
  90. }
  91. for _, h := range c.requestHandlers {
  92. if hErr := h.BeginRequest(ctx, req); hErr != nil {
  93. return hErr
  94. }
  95. }
  96. return nil
  97. }
  98. func (c *Client) emitEndRequest(ctx context.Context, resp *http.Response, err error) error {
  99. if len(c.requestHandlers) == 0 {
  100. return nil
  101. }
  102. for _, h := range c.requestHandlers {
  103. if hErr := h.EndRequest(ctx, resp, err); hErr != nil {
  104. return hErr
  105. }
  106. }
  107. return err
  108. }
  109. // RequestOption declares the type of option one can pass
  110. // to the Do methods(JSON, Form, ReadJSON...).
  111. // Request options run before request constructed.
  112. type RequestOption = func(*http.Request) error
  113. // We always add the following request headers, unless they're removed by custom ones.
  114. var defaultRequestOptions = []RequestOption{
  115. RequestHeader(false, acceptKey, contentTypeJSON),
  116. }
  117. // RequestHeader adds or sets (if overridePrev is true) a header to the request.
  118. func RequestHeader(overridePrev bool, key string, values ...string) RequestOption {
  119. key = http.CanonicalHeaderKey(key)
  120. return func(req *http.Request) error {
  121. if overridePrev { // upsert.
  122. req.Header[key] = values
  123. } else { // just insert.
  124. req.Header[key] = append(req.Header[key], values...)
  125. }
  126. return nil
  127. }
  128. }
  129. // RequestAuthorization sets an Authorization request header.
  130. // Note that we could do the same with a Transport RoundDrip too.
  131. func RequestAuthorization(value string) RequestOption {
  132. return RequestHeader(true, "Authorization", value)
  133. }
  134. // RequestAuthorizationBearer sets an Authorization: Bearer $token request header.
  135. func RequestAuthorizationBearer(accessToken string) RequestOption {
  136. headerValue := "Bearer " + accessToken
  137. return RequestAuthorization(headerValue)
  138. }
  139. // RequestQuery adds a set of URL query parameters to the request.
  140. func RequestQuery(query url.Values) RequestOption {
  141. return func(req *http.Request) error {
  142. q := req.URL.Query()
  143. for k, v := range query {
  144. q[k] = v
  145. }
  146. req.URL.RawQuery = q.Encode()
  147. return nil
  148. }
  149. }
  150. // RequestParam sets a single URL query parameter to the request.
  151. func RequestParam(key string, values ...string) RequestOption {
  152. return RequestQuery(url.Values{
  153. key: values,
  154. })
  155. }
  156. // Do sends an HTTP request and returns an HTTP response.
  157. //
  158. // The payload can be:
  159. // - io.Reader
  160. // - raw []byte
  161. // - JSON raw message
  162. // - string
  163. // - struct (JSON).
  164. //
  165. // If method is empty then it defaults to "GET".
  166. // The final variadic, optional input argument sets
  167. // the custom request options to use before the request.
  168. //
  169. // Any HTTP returned error will be of type APIError
  170. // or a timeout error if the given context was canceled.
  171. func (c *Client) Do(ctx context.Context, method, urlpath string, payload interface{}, opts ...RequestOption) (*http.Response, error) {
  172. if ctx == nil {
  173. ctx = context.Background()
  174. }
  175. if c.rateLimiter != nil {
  176. if err := c.rateLimiter.Wait(ctx); err != nil {
  177. return nil, err
  178. }
  179. }
  180. // Method defaults to GET.
  181. if method == "" {
  182. method = http.MethodGet
  183. }
  184. // Find the payload, if any.
  185. var body io.Reader
  186. if payload != nil {
  187. switch v := payload.(type) {
  188. case io.Reader:
  189. body = v
  190. case []byte:
  191. body = bytes.NewBuffer(v)
  192. case json.RawMessage:
  193. body = bytes.NewBuffer(v)
  194. case string:
  195. body = strings.NewReader(v)
  196. case url.Values:
  197. body = strings.NewReader(v.Encode())
  198. default:
  199. w := new(bytes.Buffer)
  200. // We assume it's a struct, we wont make use of reflection to find out though.
  201. err := json.NewEncoder(w).Encode(v)
  202. if err != nil {
  203. return nil, err
  204. }
  205. body = w
  206. }
  207. }
  208. if c.BaseURL != "" {
  209. urlpath = c.BaseURL + urlpath // note that we don't do any special checks here, the caller is responsible.
  210. }
  211. // Initialize the request.
  212. req, err := http.NewRequestWithContext(ctx, method, urlpath, body)
  213. if err != nil {
  214. return nil, err
  215. }
  216. // We separate the error for the default options for now.
  217. for i, opt := range c.PersistentRequestOptions {
  218. if opt == nil {
  219. continue
  220. }
  221. if err = opt(req); err != nil {
  222. return nil, fmt.Errorf("client.Do: default request option[%d]: %w", i, err)
  223. }
  224. }
  225. // Apply any custom request options (e.g. content type, accept headers, query...)
  226. for _, opt := range opts {
  227. if opt == nil {
  228. continue
  229. }
  230. if err = opt(req); err != nil {
  231. return nil, err
  232. }
  233. }
  234. if err = c.emitBeginRequest(ctx, req); err != nil {
  235. return nil, err
  236. }
  237. // Caller is responsible for closing the response body.
  238. // Also note that the gzip compression is handled automatically nowadays.
  239. resp, respErr := c.HTTPClient.Do(req)
  240. if err = c.emitEndRequest(ctx, resp, respErr); err != nil {
  241. return nil, err
  242. }
  243. return resp, respErr
  244. }
  245. // DrainResponseBody drains response body and close it, allowing the transport to reuse TCP connections.
  246. // It's automatically called on Client.ReadXXX methods on the end.
  247. func (c *Client) DrainResponseBody(resp *http.Response) {
  248. _, _ = io.Copy(io.Discard, resp.Body)
  249. resp.Body.Close()
  250. }
  251. const (
  252. acceptKey = "Accept"
  253. contentTypeKey = "Content-Type"
  254. contentLengthKey = "Content-Length"
  255. contentTypePlainText = "plain/text"
  256. contentTypeJSON = "application/json"
  257. contentTypeFormURLEncoded = "application/x-www-form-urlencoded"
  258. )
  259. // JSON writes data as JSON to the server.
  260. func (c *Client) JSON(ctx context.Context, method, urlpath string, payload interface{}, opts ...RequestOption) (*http.Response, error) {
  261. opts = append(opts, RequestHeader(true, contentTypeKey, contentTypeJSON))
  262. return c.Do(ctx, method, urlpath, payload, opts...)
  263. }
  264. // JSON writes form data to the server.
  265. func (c *Client) Form(ctx context.Context, method, urlpath string, formValues url.Values, opts ...RequestOption) (*http.Response, error) {
  266. payload := formValues.Encode()
  267. opts = append(opts,
  268. RequestHeader(true, contentTypeKey, contentTypeFormURLEncoded),
  269. RequestHeader(true, contentLengthKey, strconv.Itoa(len(payload))),
  270. )
  271. return c.Do(ctx, method, urlpath, payload, opts...)
  272. }
  273. // Uploader holds the necessary information for upload requests.
  274. //
  275. // Look the Client.NewUploader method.
  276. type Uploader struct {
  277. client *Client
  278. body *bytes.Buffer
  279. Writer *multipart.Writer
  280. }
  281. // AddFileSource adds a form field to the uploader with the given key.
  282. func (u *Uploader) AddField(key, value string) error {
  283. f, err := u.Writer.CreateFormField(key)
  284. if err != nil {
  285. return err
  286. }
  287. _, err = io.Copy(f, strings.NewReader(value))
  288. return err
  289. }
  290. // AddFileSource adds a form file to the uploader with the given key.
  291. func (u *Uploader) AddFileSource(key, filename string, source io.Reader) error {
  292. f, err := u.Writer.CreateFormFile(key, filename)
  293. if err != nil {
  294. return err
  295. }
  296. _, err = io.Copy(f, source)
  297. return err
  298. }
  299. // AddFile adds a local form file to the uploader with the given key.
  300. func (u *Uploader) AddFile(key, filename string) error {
  301. source, err := os.Open(filename)
  302. if err != nil {
  303. return err
  304. }
  305. return u.AddFileSource(key, filename, source)
  306. }
  307. // Uploads sends local data to the server.
  308. func (u *Uploader) Upload(ctx context.Context, method, urlpath string, opts ...RequestOption) (*http.Response, error) {
  309. err := u.Writer.Close()
  310. if err != nil {
  311. return nil, err
  312. }
  313. payload := bytes.NewReader(u.body.Bytes())
  314. opts = append(opts, RequestHeader(true, contentTypeKey, u.Writer.FormDataContentType()))
  315. return u.client.Do(ctx, method, urlpath, payload, opts...)
  316. }
  317. // NewUploader returns a structure which is responsible for sending
  318. // file and form data to the server.
  319. func (c *Client) NewUploader() *Uploader {
  320. body := new(bytes.Buffer)
  321. writer := multipart.NewWriter(body)
  322. return &Uploader{
  323. client: c,
  324. body: body,
  325. Writer: writer,
  326. }
  327. }
  328. // ReadJSON binds "dest" to the response's body.
  329. // After this call, the response body reader is closed.
  330. func (c *Client) ReadJSON(ctx context.Context, dest interface{}, method, urlpath string, payload interface{}, opts ...RequestOption) error {
  331. if payload != nil {
  332. opts = append(opts, RequestHeader(true, contentTypeKey, contentTypeJSON))
  333. }
  334. resp, err := c.Do(ctx, method, urlpath, payload, opts...)
  335. if err != nil {
  336. return err
  337. }
  338. defer c.DrainResponseBody(resp)
  339. if resp.StatusCode >= http.StatusBadRequest {
  340. return ExtractError(resp)
  341. }
  342. // DBUG
  343. // b, _ := io.ReadAll(resp.Body)
  344. // println(string(b))
  345. // return json.Unmarshal(b, &dest)
  346. if dest != nil {
  347. return json.NewDecoder(resp.Body).Decode(&dest)
  348. }
  349. return json.NewDecoder(resp.Body).Decode(&dest)
  350. }
  351. // ReadPlain like ReadJSON but it accepts a pointer to a string or byte slice or integer
  352. // and it reads the body as plain text.
  353. func (c *Client) ReadPlain(ctx context.Context, dest interface{}, method, urlpath string, payload interface{}, opts ...RequestOption) error {
  354. resp, err := c.Do(ctx, method, urlpath, payload, opts...)
  355. if err != nil {
  356. return err
  357. }
  358. defer c.DrainResponseBody(resp)
  359. if resp.StatusCode >= http.StatusBadRequest {
  360. return ExtractError(resp)
  361. }
  362. body, err := io.ReadAll(resp.Body)
  363. if err != nil {
  364. return err
  365. }
  366. switch ptr := dest.(type) {
  367. case *[]byte:
  368. *ptr = body
  369. return nil
  370. case *string:
  371. *ptr = string(body)
  372. return nil
  373. case *int:
  374. *ptr, err = strconv.Atoi(string(body))
  375. return err
  376. default:
  377. return fmt.Errorf("unsupported response body type: %T", ptr)
  378. }
  379. }
  380. // GetPlainUnquote reads the response body as raw text and tries to unquote it,
  381. // useful when the remote server sends a single key as a value but due to backend mistake
  382. // it sends it as JSON (quoted) instead of plain text.
  383. func (c *Client) GetPlainUnquote(ctx context.Context, method, urlpath string, payload interface{}, opts ...RequestOption) (string, error) {
  384. var bodyStr string
  385. if err := c.ReadPlain(ctx, &bodyStr, method, urlpath, payload, opts...); err != nil {
  386. return "", err
  387. }
  388. s, err := strconv.Unquote(bodyStr)
  389. if err == nil {
  390. bodyStr = s
  391. }
  392. return bodyStr, nil
  393. }
  394. // WriteTo reads the response and then copies its data to the "dest" writer.
  395. // If the "dest" is a type of HTTP response writer then it writes the
  396. // content-type and content-length of the original request.
  397. //
  398. // Returns the amount of bytes written to "dest".
  399. func (c *Client) WriteTo(ctx context.Context, dest io.Writer, method, urlpath string, payload interface{}, opts ...RequestOption) (int64, error) {
  400. if payload != nil {
  401. opts = append(opts, RequestHeader(true, contentTypeKey, contentTypeJSON))
  402. }
  403. resp, err := c.Do(ctx, method, urlpath, payload, opts...)
  404. if err != nil {
  405. return 0, err
  406. }
  407. defer resp.Body.Close()
  408. if w, ok := dest.(http.ResponseWriter); ok {
  409. // Copy the content type and content-length.
  410. w.Header().Set("Content-Type", resp.Header.Get("Content-Type"))
  411. if resp.ContentLength > 0 {
  412. w.Header().Set("Content-Length", strconv.FormatInt(resp.ContentLength, 10))
  413. }
  414. }
  415. return io.Copy(dest, resp.Body)
  416. }
  417. // BindResponse consumes the response's body and binds the result to the "dest" pointer,
  418. // closing the response's body is up to the caller.
  419. //
  420. // The "dest" will be binded based on the response's content type header.
  421. // Note that this is strict in order to catch bad actioners fast,
  422. // e.g. it wont try to read plain text if not specified on
  423. // the response headers and the dest is a *string.
  424. func BindResponse(resp *http.Response, dest interface{}) (err error) {
  425. contentType := trimHeader(resp.Header.Get(contentTypeKey))
  426. switch contentType {
  427. case contentTypeJSON: // the most common scenario on successful responses.
  428. return json.NewDecoder(resp.Body).Decode(&dest)
  429. case contentTypePlainText:
  430. b, err := io.ReadAll(resp.Body)
  431. if err != nil {
  432. return err
  433. }
  434. switch v := dest.(type) {
  435. case *string:
  436. *v = string(b)
  437. case *[]byte:
  438. *v = b
  439. default:
  440. return fmt.Errorf("plain text response should accept a *string or a *[]byte")
  441. }
  442. default:
  443. acceptContentType := trimHeader(resp.Request.Header.Get(acceptKey))
  444. msg := ""
  445. if acceptContentType == contentType {
  446. // Here we make a special case, if the content type
  447. // was explicitly set by the request but we cannot handle it.
  448. msg = fmt.Sprintf("current implementation can not handle the received (and accepted) mime type: %s", contentType)
  449. } else {
  450. msg = fmt.Sprintf("unexpected mime type received: %s", contentType)
  451. }
  452. err = errors.New(msg)
  453. }
  454. return
  455. }
  456. func trimHeader(v string) string {
  457. for i, char := range v {
  458. if char == ' ' || char == ';' {
  459. return v[:i]
  460. }
  461. }
  462. return v
  463. }