Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: auto parse filename from content-disposition or URL #926 #932

Merged
merged 1 commit into from
Dec 30, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 29 additions & 2 deletions client.go
Original file line number Diff line number Diff line change
Expand Up @@ -201,6 +201,7 @@ type Client struct {
isTrace bool
debugBodyLimit int
outputDirectory string
isSaveResponse bool
scheme string
log Logger
ctx context.Context
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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 {
Expand All @@ -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) {
Expand Down Expand Up @@ -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]
Expand Down
29 changes: 23 additions & 6 deletions middleware.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,10 +9,12 @@ import (
"bytes"
"fmt"
"io"
"mime"
"mime/multipart"
"net/http"
"net/textproto"
"net/url"
"path"
"path/filepath"
"reflect"
"strconv"
Expand Down Expand Up @@ -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
}
Expand Down
6 changes: 3 additions & 3 deletions middleware_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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())
}
Expand All @@ -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())
}
Expand Down
31 changes: 23 additions & 8 deletions request.go
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ type Request struct {
Debug bool
CloseConnection bool
DoNotParseResponse bool
OutputFile string
OutputFileName string
ExpectResponseContentType string
ForceResponseContentType string
DebugBodyLimit int
Expand All @@ -60,6 +60,7 @@ type Request struct {
AllowMethodGetPayload bool
AllowMethodDeletePayload bool
IsDone bool
IsSaveResponse bool
Timeout time.Duration
RetryCount int
RetryWaitTime time.Duration
Expand All @@ -81,7 +82,6 @@ type Request struct {
isMultiPart bool
isFormData bool
setContentLength bool
isSaveResponse bool
jsonEscapeHTML bool
ctx context.Context
ctxCancelFunc context.CancelFunc
Expand Down Expand Up @@ -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.
Expand All @@ -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
}

Expand Down Expand Up @@ -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]
Expand Down
54 changes: 49 additions & 5 deletions request_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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()
Expand Down Expand Up @@ -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)
Expand Down
4 changes: 4 additions & 0 deletions resty_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion util.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading