From fdf601a4a56d1877af3d60992a0cec78205accbb Mon Sep 17 00:00:00 2001 From: "Jeevanandam M." Date: Sun, 29 Dec 2024 20:14:11 -0800 Subject: [PATCH] feat: auto parse filename from content-disposition or URL #926 (#932) --- client.go | 31 ++++++++++++++++++++++++-- middleware.go | 29 +++++++++++++++++++------ middleware_test.go | 6 +++--- request.go | 31 +++++++++++++++++++------- request_test.go | 54 +++++++++++++++++++++++++++++++++++++++++----- resty_test.go | 4 ++++ util.go | 2 +- 7 files changed, 132 insertions(+), 25 deletions(-) diff --git a/client.go b/client.go index 909b07cc..0d089f69 100644 --- a/client.go +++ b/client.go @@ -201,6 +201,7 @@ type Client struct { isTrace bool debugBodyLimit int outputDirectory string + isSaveResponse bool scheme string log Logger ctx context.Context @@ -616,6 +617,7 @@ func (c *Client) R() *Request { Timeout: c.timeout, Debug: c.debug, IsTrace: c.isTrace, + IsSaveResponse: c.isSaveResponse, AuthScheme: c.authScheme, AuthToken: c.authToken, RetryCount: c.retryCount, @@ -1646,7 +1648,7 @@ func (c *Client) OutputDirectory() string { // SetOutputDirectory method sets the output directory for saving HTTP responses in a file. // Resty creates one if the output directory does not exist. This setting is optional, -// if you plan to use the absolute path in [Request.SetOutputFile] and can used together. +// if you plan to use the absolute path in [Request.SetOutputFileName] and can used together. // // client.SetOutputDirectory("/save/http/response/here") func (c *Client) SetOutputDirectory(dirPath string) *Client { @@ -1656,6 +1658,31 @@ func (c *Client) SetOutputDirectory(dirPath string) *Client { return c } +// IsSaveResponse method returns true if the save response is set to true; otherwise, false +func (c *Client) IsSaveResponse() bool { + c.lock.RLock() + defer c.lock.RUnlock() + return c.isSaveResponse +} + +// SetSaveResponse method used to enable the save response option at the client level for +// all requests +// +// client.SetSaveResponse(true) +// +// Resty determines the save filename in the following order - +// - [Request.SetOutputFileName] +// - Content-Disposition header +// - Request URL using [path.Base] +// +// It can be overridden at request level, see [Request.SetSaveResponse] +func (c *Client) SetSaveResponse(save bool) *Client { + c.lock.Lock() + defer c.lock.Unlock() + c.isSaveResponse = save + return c +} + // HTTPTransport method does type assertion and returns [http.Transport] // from the client instance, if type assertion fails it returns an error func (c *Client) HTTPTransport() (*http.Transport, error) { @@ -1875,7 +1902,7 @@ func (c *Client) ResponseBodyLimit() int64 { // in the uncompressed response is larger than the limit. // Body size limit will not be enforced in the following cases: // - ResponseBodyLimit <= 0, which is the default behavior. -// - [Request.SetOutputFile] is called to save response data to the file. +// - [Request.SetOutputFileName] is called to save response data to the file. // - "DoNotParseResponse" is set for client or request. // // It can be overridden at the request level; see [Request.SetResponseBodyLimit] diff --git a/middleware.go b/middleware.go index 94be3b8b..98557f6c 100644 --- a/middleware.go +++ b/middleware.go @@ -9,10 +9,12 @@ import ( "bytes" "fmt" "io" + "mime" "mime/multipart" "net/http" "net/textproto" "net/url" + "path" "path/filepath" "reflect" "strconv" @@ -547,19 +549,34 @@ func AutoParseResponseMiddleware(c *Client, res *Response) (err error) { } // SaveToFileResponseMiddleware method used to write HTTP response body into -// given file details via [Request.SetOutputFile] +// file. The filename is determined in the following order - +// - [Request.SetOutputFileName] +// - Content-Disposition header +// - Request URL using [path.Base] func SaveToFileResponseMiddleware(c *Client, res *Response) error { - if res.Err != nil || !res.Request.isSaveResponse { + if res.Err != nil || !res.Request.IsSaveResponse { return nil } - file := "" + file := res.Request.OutputFileName + if isStringEmpty(file) { + cntDispositionValue := res.Header().Get(hdrContentDisposition) + if len(cntDispositionValue) > 0 { + if _, params, err := mime.ParseMediaType(cntDispositionValue); err == nil { + file = params["filename"] + } + } + if isStringEmpty(file) { + urlPath, _ := url.Parse(res.Request.URL) + file = path.Base(urlPath.Path) + } + } - if len(c.OutputDirectory()) > 0 && !filepath.IsAbs(res.Request.OutputFile) { - file += c.OutputDirectory() + string(filepath.Separator) + if len(c.OutputDirectory()) > 0 && !filepath.IsAbs(file) { + file = filepath.Join(c.OutputDirectory(), string(filepath.Separator), file) } - file = filepath.Clean(file + res.Request.OutputFile) + file = filepath.Clean(file) if err := createDirectory(filepath.Dir(file)); err != nil { return err } diff --git a/middleware_test.go b/middleware_test.go index d11ba215..29c67f34 100644 --- a/middleware_test.go +++ b/middleware_test.go @@ -878,13 +878,13 @@ func TestMiddlewareSaveToFileErrorCases(t *testing.T) { // dir create error req1 := c.R() - req1.SetOutputFile(filepath.Join(tempDir, "new-res-dir", "sample.txt")) + req1.SetOutputFileName(filepath.Join(tempDir, "new-res-dir", "sample.txt")) err1 := SaveToFileResponseMiddleware(c, &Response{Request: req1}) assertEqual(t, errDirMsg, err1.Error()) // file create error req2 := c.R() - req2.SetOutputFile(filepath.Join(tempDir, "sample.txt")) + req2.SetOutputFileName(filepath.Join(tempDir, "sample.txt")) err2 := SaveToFileResponseMiddleware(c, &Response{Request: req2}) assertEqual(t, errFileMsg, err2.Error()) } @@ -903,7 +903,7 @@ func TestMiddlewareSaveToFileCopyError(t *testing.T) { // copy error req1 := c.R() - req1.SetOutputFile(filepath.Join(tempDir, "new-res-dir", "sample.txt")) + req1.SetOutputFileName(filepath.Join(tempDir, "new-res-dir", "sample.txt")) err1 := SaveToFileResponseMiddleware(c, &Response{Request: req1, Body: io.NopCloser(bytes.NewBufferString("Test context"))}) assertEqual(t, errCopyMsg, err1.Error()) } diff --git a/request.go b/request.go index 4e8cfe15..21e211f6 100644 --- a/request.go +++ b/request.go @@ -50,7 +50,7 @@ type Request struct { Debug bool CloseConnection bool DoNotParseResponse bool - OutputFile string + OutputFileName string ExpectResponseContentType string ForceResponseContentType string DebugBodyLimit int @@ -60,6 +60,7 @@ type Request struct { AllowMethodGetPayload bool AllowMethodDeletePayload bool IsDone bool + IsSaveResponse bool Timeout time.Duration RetryCount int RetryWaitTime time.Duration @@ -81,7 +82,6 @@ type Request struct { isMultiPart bool isFormData bool setContentLength bool - isSaveResponse bool jsonEscapeHTML bool ctx context.Context ctxCancelFunc context.CancelFunc @@ -662,7 +662,7 @@ func (r *Request) SetAuthScheme(scheme string) *Request { return r } -// SetOutputFile method sets the output file for the current HTTP request. The current +// SetOutputFileName method sets the output file for the current HTTP request. The current // HTTP response will be saved in the given file. It is similar to the `curl -o` flag. // // Absolute path or relative path can be used. @@ -671,15 +671,30 @@ func (r *Request) SetAuthScheme(scheme string) *Request { // in the [Client.SetOutputDirectory]. // // client.R(). -// SetOutputFile("/Users/jeeva/Downloads/ReplyWithHeader-v5.1-beta.zip"). +// SetOutputFileName("/Users/jeeva/Downloads/ReplyWithHeader-v5.1-beta.zip"). // Get("http://bit.ly/1LouEKr") // // NOTE: In this scenario // - [Response.BodyBytes] might be nil. // - [Response].Body might be already read. -func (r *Request) SetOutputFile(file string) *Request { - r.OutputFile = file - r.isSaveResponse = true +func (r *Request) SetOutputFileName(file string) *Request { + r.OutputFileName = file + r.SetSaveResponse(true) + return r +} + +// SetSaveResponse method used to enable the save response option for the current requests +// +// client.R().SetSaveResponse(true) +// +// Resty determines the save filename in the following order - +// - [Request.SetOutputFileName] +// - Content-Disposition header +// - Request URL using [path.Base] +// +// It overrides the value set at the client instance level, see [Client.SetSaveResponse] +func (r *Request) SetSaveResponse(save bool) *Request { + r.IsSaveResponse = save return r } @@ -711,7 +726,7 @@ func (r *Request) SetDoNotParseResponse(notParse bool) *Request { // in the uncompressed response is larger than the limit. // Body size limit will not be enforced in the following cases: // - ResponseBodyLimit <= 0, which is the default behavior. -// - [Request.SetOutputFile] is called to save response data to the file. +// - [Request.SetOutputFileName] is called to save response data to the file. // - "DoNotParseResponse" is set for client or request. // // It overrides the value set at the client instance level, see [Client.SetResponseBodyLimit] diff --git a/request_test.go b/request_test.go index 60e8e2c2..091c446d 100644 --- a/request_test.go +++ b/request_test.go @@ -1299,11 +1299,10 @@ func TestOutputFileWithBaseDirAndRelativePath(t *testing.T) { SetRedirectPolicy(FlexibleRedirectPolicy(10)). SetOutputDirectory(baseOutputDir). SetDebug(true) - client.outputLogTo(io.Discard) outputFilePath := "go-resty/test-img-success.png" resp, err := client.R(). - SetOutputFile(outputFilePath). + SetOutputFileName(outputFilePath). Get(ts.URL + "/my-image.png") assertError(t, err) @@ -1332,7 +1331,7 @@ func TestOutputPathDirNotExists(t *testing.T) { SetOutputDirectory(filepath.Join(getTestDataPath(), "not-exists-dir")) resp, err := client.R(). - SetOutputFile("test-img-success.png"). + SetOutputFileName("test-img-success.png"). Get(ts.URL + "/my-image.png") assertError(t, err) @@ -1348,7 +1347,7 @@ func TestOutputFileAbsPath(t *testing.T) { outputFile := filepath.Join(getTestDataPath(), "go-resty", "test-img-success-2.png") res, err := dcnlr(). - SetOutputFile(outputFile). + SetOutputFileName(outputFile). Get(ts.URL + "/my-image.png") assertError(t, err) @@ -1358,6 +1357,51 @@ func TestOutputFileAbsPath(t *testing.T) { assertNil(t, err) } +func TestRequestSaveResponse(t *testing.T) { + ts := createGetServer(t) + defer ts.Close() + defer cleanupFiles(filepath.Join(".testdata", "go-resty")) + + c := dcnl(). + SetSaveResponse(true). + SetOutputDirectory(filepath.Join(getTestDataPath(), "go-resty")) + + assertEqual(t, true, c.IsSaveResponse()) + + t.Run("content-disposition save response request", func(t *testing.T) { + outputFile := filepath.Join(getTestDataPath(), "go-resty", "test-img-success-2.png") + c.SetSaveResponse(false) + assertEqual(t, false, c.IsSaveResponse()) + + res, err := c.R(). + SetSaveResponse(true). + Get(ts.URL + "/my-image.png?content-disposition=true&filename=test-img-success-2.png") + + assertError(t, err) + assertEqual(t, int64(2579468), res.Size()) + + _, err = os.Stat(outputFile) + assertNil(t, err) + }) + + t.Run("use filename from path", func(t *testing.T) { + outputFile := filepath.Join(getTestDataPath(), "go-resty", "my-image.png") + c.SetSaveResponse(false) + assertEqual(t, false, c.IsSaveResponse()) + + res, err := c.R(). + SetSaveResponse(true). + Get(ts.URL + "/my-image.png") + + assertError(t, err) + assertEqual(t, int64(2579468), res.Size()) + + _, err = os.Stat(outputFile) + assertNil(t, err) + }) + +} + func TestContextInternal(t *testing.T) { ts := createGetServer(t) defer ts.Close() @@ -2175,7 +2219,7 @@ func TestRequestSetResultAndSetOutputFile(t *testing.T) { SetBody(&credentials{Username: "testuser", Password: "testpass"}). SetResponseBodyUnlimitedReads(true). SetResult(&AuthSuccess{}). - SetOutputFile(outputFile). + SetOutputFileName(outputFile). Post("/login") assertError(t, err) diff --git a/resty_test.go b/resty_test.go index bf7084a8..deb4aea7 100644 --- a/resty_test.go +++ b/resty_test.go @@ -114,6 +114,10 @@ func createGetServer(t *testing.T) *httptest.Server { fileBytes, _ := os.ReadFile(filepath.Join(getTestDataPath(), "test-img.png")) w.Header().Set("Content-Type", "image/png") w.Header().Set("Content-Length", strconv.Itoa(len(fileBytes))) + if r.URL.Query().Get("content-disposition") == "true" { + filename := r.URL.Query().Get("filename") + w.Header().Set(hdrContentDisposition, "inline; filename=\""+filename+"\"") + } _, _ = w.Write(fileBytes) case "/get-method-payload-test": body, err := io.ReadAll(r.Body) diff --git a/util.go b/util.go index 942bd9ba..bba6c04c 100644 --- a/util.go +++ b/util.go @@ -407,7 +407,7 @@ func responseDebugLogger(c *Client, res *Response) { fmt.Sprintf("TIME DURATION: %v\n", res.Time()) + "HEADERS :\n" + composeHeaders(rl.Header) + "\n" - if res.Request.isSaveResponse { + if res.Request.IsSaveResponse { debugLog += "BODY :\n***** RESPONSE WRITTEN INTO FILE *****\n" } else { debugLog += fmt.Sprintf("BODY :\n%v\n", rl.Body)