From ad38855261dca802439922f71408e2b08e7c10ea Mon Sep 17 00:00:00 2001 From: r Date: Sun, 17 Sep 2023 09:44:02 +0000 Subject: Set timeout and response size limit for the http client --- mastodon/apps.go | 3 +-- mastodon/http.go | 41 +++++++++++++++++++++++++++++++++++++++++ mastodon/mastodon.go | 3 ++- 3 files changed, 44 insertions(+), 3 deletions(-) create mode 100644 mastodon/http.go diff --git a/mastodon/apps.go b/mastodon/apps.go index 5d925c3..12d2e86 100644 --- a/mastodon/apps.go +++ b/mastodon/apps.go @@ -11,7 +11,6 @@ import ( // AppConfig is a setting for registering applications. type AppConfig struct { - http.Client Server string ClientName string @@ -62,7 +61,7 @@ func RegisterApp(ctx context.Context, appConfig *AppConfig) (*Application, error } req = req.WithContext(ctx) req.Header.Set("Content-Type", "application/x-www-form-urlencoded") - resp, err := appConfig.Do(req) + resp, err := httpClient.Do(req) if err != nil { return nil, err } diff --git a/mastodon/http.go b/mastodon/http.go new file mode 100644 index 0000000..ca4c771 --- /dev/null +++ b/mastodon/http.go @@ -0,0 +1,41 @@ +package mastodon + +import ( + "fmt" + "io" + "net/http" + "time" +) + +type lr struct { + io.ReadCloser + r *http.Request +} + +func (r *lr) Read(p []byte) (n int, err error) { + n, err = r.ReadCloser.Read(p) + // override the generic error returned by the MaxBytesReader + if _, ok := err.(*http.MaxBytesError); ok { + err = fmt.Errorf("%s \"%s\": response body too large", r.r.Method, r.r.URL) + } + return +} + +type transport struct { + t http.RoundTripper +} + +func (t *transport) RoundTrip(r *http.Request) (*http.Response, error) { + resp, err := t.t.RoundTrip(r) + if resp != nil && resp.Body != nil { + resp.Body = &lr{http.MaxBytesReader(nil, resp.Body, 8<<20), r} + } + return resp, err +} + +var httpClient = &http.Client{ + Transport: &transport{ + t: http.DefaultTransport, + }, + Timeout: 30 * time.Second, +} diff --git a/mastodon/mastodon.go b/mastodon/mastodon.go index 8678314..f114169 100644 --- a/mastodon/mastodon.go +++ b/mastodon/mastodon.go @@ -168,12 +168,13 @@ func (c *Client) doAPI(ctx context.Context, method string, uri string, params in } } return json.NewDecoder(resp.Body).Decode(&res) + } // NewClient return new mastodon API client. func NewClient(config *Config) *Client { return &Client{ - Client: http.DefaultClient, + Client: httpClient, config: config, } } -- cgit v1.2.3