Skip to content

Commit

Permalink
feat: auto parse filename from content-disposition or URL #926 (#932)
Browse files Browse the repository at this point in the history
  • Loading branch information
jeevatkm authored Dec 30, 2024
1 parent af72a4d commit fdf601a
Show file tree
Hide file tree
Showing 7 changed files with 132 additions and 25 deletions.
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

0 comments on commit fdf601a

Please sign in to comment.