Skip to content

Commit fdf601a

Browse files
authored
feat: auto parse filename from content-disposition or URL #926 (#932)
1 parent af72a4d commit fdf601a

File tree

7 files changed

+132
-25
lines changed

7 files changed

+132
-25
lines changed

client.go

+29-2
Original file line numberDiff line numberDiff line change
@@ -201,6 +201,7 @@ type Client struct {
201201
isTrace bool
202202
debugBodyLimit int
203203
outputDirectory string
204+
isSaveResponse bool
204205
scheme string
205206
log Logger
206207
ctx context.Context
@@ -616,6 +617,7 @@ func (c *Client) R() *Request {
616617
Timeout: c.timeout,
617618
Debug: c.debug,
618619
IsTrace: c.isTrace,
620+
IsSaveResponse: c.isSaveResponse,
619621
AuthScheme: c.authScheme,
620622
AuthToken: c.authToken,
621623
RetryCount: c.retryCount,
@@ -1646,7 +1648,7 @@ func (c *Client) OutputDirectory() string {
16461648

16471649
// SetOutputDirectory method sets the output directory for saving HTTP responses in a file.
16481650
// Resty creates one if the output directory does not exist. This setting is optional,
1649-
// if you plan to use the absolute path in [Request.SetOutputFile] and can used together.
1651+
// if you plan to use the absolute path in [Request.SetOutputFileName] and can used together.
16501652
//
16511653
// client.SetOutputDirectory("/save/http/response/here")
16521654
func (c *Client) SetOutputDirectory(dirPath string) *Client {
@@ -1656,6 +1658,31 @@ func (c *Client) SetOutputDirectory(dirPath string) *Client {
16561658
return c
16571659
}
16581660

1661+
// IsSaveResponse method returns true if the save response is set to true; otherwise, false
1662+
func (c *Client) IsSaveResponse() bool {
1663+
c.lock.RLock()
1664+
defer c.lock.RUnlock()
1665+
return c.isSaveResponse
1666+
}
1667+
1668+
// SetSaveResponse method used to enable the save response option at the client level for
1669+
// all requests
1670+
//
1671+
// client.SetSaveResponse(true)
1672+
//
1673+
// Resty determines the save filename in the following order -
1674+
// - [Request.SetOutputFileName]
1675+
// - Content-Disposition header
1676+
// - Request URL using [path.Base]
1677+
//
1678+
// It can be overridden at request level, see [Request.SetSaveResponse]
1679+
func (c *Client) SetSaveResponse(save bool) *Client {
1680+
c.lock.Lock()
1681+
defer c.lock.Unlock()
1682+
c.isSaveResponse = save
1683+
return c
1684+
}
1685+
16591686
// HTTPTransport method does type assertion and returns [http.Transport]
16601687
// from the client instance, if type assertion fails it returns an error
16611688
func (c *Client) HTTPTransport() (*http.Transport, error) {
@@ -1875,7 +1902,7 @@ func (c *Client) ResponseBodyLimit() int64 {
18751902
// in the uncompressed response is larger than the limit.
18761903
// Body size limit will not be enforced in the following cases:
18771904
// - ResponseBodyLimit <= 0, which is the default behavior.
1878-
// - [Request.SetOutputFile] is called to save response data to the file.
1905+
// - [Request.SetOutputFileName] is called to save response data to the file.
18791906
// - "DoNotParseResponse" is set for client or request.
18801907
//
18811908
// It can be overridden at the request level; see [Request.SetResponseBodyLimit]

middleware.go

+23-6
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,12 @@ import (
99
"bytes"
1010
"fmt"
1111
"io"
12+
"mime"
1213
"mime/multipart"
1314
"net/http"
1415
"net/textproto"
1516
"net/url"
17+
"path"
1618
"path/filepath"
1719
"reflect"
1820
"strconv"
@@ -547,19 +549,34 @@ func AutoParseResponseMiddleware(c *Client, res *Response) (err error) {
547549
}
548550

549551
// SaveToFileResponseMiddleware method used to write HTTP response body into
550-
// given file details via [Request.SetOutputFile]
552+
// file. The filename is determined in the following order -
553+
// - [Request.SetOutputFileName]
554+
// - Content-Disposition header
555+
// - Request URL using [path.Base]
551556
func SaveToFileResponseMiddleware(c *Client, res *Response) error {
552-
if res.Err != nil || !res.Request.isSaveResponse {
557+
if res.Err != nil || !res.Request.IsSaveResponse {
553558
return nil
554559
}
555560

556-
file := ""
561+
file := res.Request.OutputFileName
562+
if isStringEmpty(file) {
563+
cntDispositionValue := res.Header().Get(hdrContentDisposition)
564+
if len(cntDispositionValue) > 0 {
565+
if _, params, err := mime.ParseMediaType(cntDispositionValue); err == nil {
566+
file = params["filename"]
567+
}
568+
}
569+
if isStringEmpty(file) {
570+
urlPath, _ := url.Parse(res.Request.URL)
571+
file = path.Base(urlPath.Path)
572+
}
573+
}
557574

558-
if len(c.OutputDirectory()) > 0 && !filepath.IsAbs(res.Request.OutputFile) {
559-
file += c.OutputDirectory() + string(filepath.Separator)
575+
if len(c.OutputDirectory()) > 0 && !filepath.IsAbs(file) {
576+
file = filepath.Join(c.OutputDirectory(), string(filepath.Separator), file)
560577
}
561578

562-
file = filepath.Clean(file + res.Request.OutputFile)
579+
file = filepath.Clean(file)
563580
if err := createDirectory(filepath.Dir(file)); err != nil {
564581
return err
565582
}

middleware_test.go

+3-3
Original file line numberDiff line numberDiff line change
@@ -878,13 +878,13 @@ func TestMiddlewareSaveToFileErrorCases(t *testing.T) {
878878

879879
// dir create error
880880
req1 := c.R()
881-
req1.SetOutputFile(filepath.Join(tempDir, "new-res-dir", "sample.txt"))
881+
req1.SetOutputFileName(filepath.Join(tempDir, "new-res-dir", "sample.txt"))
882882
err1 := SaveToFileResponseMiddleware(c, &Response{Request: req1})
883883
assertEqual(t, errDirMsg, err1.Error())
884884

885885
// file create error
886886
req2 := c.R()
887-
req2.SetOutputFile(filepath.Join(tempDir, "sample.txt"))
887+
req2.SetOutputFileName(filepath.Join(tempDir, "sample.txt"))
888888
err2 := SaveToFileResponseMiddleware(c, &Response{Request: req2})
889889
assertEqual(t, errFileMsg, err2.Error())
890890
}
@@ -903,7 +903,7 @@ func TestMiddlewareSaveToFileCopyError(t *testing.T) {
903903

904904
// copy error
905905
req1 := c.R()
906-
req1.SetOutputFile(filepath.Join(tempDir, "new-res-dir", "sample.txt"))
906+
req1.SetOutputFileName(filepath.Join(tempDir, "new-res-dir", "sample.txt"))
907907
err1 := SaveToFileResponseMiddleware(c, &Response{Request: req1, Body: io.NopCloser(bytes.NewBufferString("Test context"))})
908908
assertEqual(t, errCopyMsg, err1.Error())
909909
}

request.go

+23-8
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ type Request struct {
5050
Debug bool
5151
CloseConnection bool
5252
DoNotParseResponse bool
53-
OutputFile string
53+
OutputFileName string
5454
ExpectResponseContentType string
5555
ForceResponseContentType string
5656
DebugBodyLimit int
@@ -60,6 +60,7 @@ type Request struct {
6060
AllowMethodGetPayload bool
6161
AllowMethodDeletePayload bool
6262
IsDone bool
63+
IsSaveResponse bool
6364
Timeout time.Duration
6465
RetryCount int
6566
RetryWaitTime time.Duration
@@ -81,7 +82,6 @@ type Request struct {
8182
isMultiPart bool
8283
isFormData bool
8384
setContentLength bool
84-
isSaveResponse bool
8585
jsonEscapeHTML bool
8686
ctx context.Context
8787
ctxCancelFunc context.CancelFunc
@@ -662,7 +662,7 @@ func (r *Request) SetAuthScheme(scheme string) *Request {
662662
return r
663663
}
664664

665-
// SetOutputFile method sets the output file for the current HTTP request. The current
665+
// SetOutputFileName method sets the output file for the current HTTP request. The current
666666
// HTTP response will be saved in the given file. It is similar to the `curl -o` flag.
667667
//
668668
// Absolute path or relative path can be used.
@@ -671,15 +671,30 @@ func (r *Request) SetAuthScheme(scheme string) *Request {
671671
// in the [Client.SetOutputDirectory].
672672
//
673673
// client.R().
674-
// SetOutputFile("/Users/jeeva/Downloads/ReplyWithHeader-v5.1-beta.zip").
674+
// SetOutputFileName("/Users/jeeva/Downloads/ReplyWithHeader-v5.1-beta.zip").
675675
// Get("http://bit.ly/1LouEKr")
676676
//
677677
// NOTE: In this scenario
678678
// - [Response.BodyBytes] might be nil.
679679
// - [Response].Body might be already read.
680-
func (r *Request) SetOutputFile(file string) *Request {
681-
r.OutputFile = file
682-
r.isSaveResponse = true
680+
func (r *Request) SetOutputFileName(file string) *Request {
681+
r.OutputFileName = file
682+
r.SetSaveResponse(true)
683+
return r
684+
}
685+
686+
// SetSaveResponse method used to enable the save response option for the current requests
687+
//
688+
// client.R().SetSaveResponse(true)
689+
//
690+
// Resty determines the save filename in the following order -
691+
// - [Request.SetOutputFileName]
692+
// - Content-Disposition header
693+
// - Request URL using [path.Base]
694+
//
695+
// It overrides the value set at the client instance level, see [Client.SetSaveResponse]
696+
func (r *Request) SetSaveResponse(save bool) *Request {
697+
r.IsSaveResponse = save
683698
return r
684699
}
685700

@@ -711,7 +726,7 @@ func (r *Request) SetDoNotParseResponse(notParse bool) *Request {
711726
// in the uncompressed response is larger than the limit.
712727
// Body size limit will not be enforced in the following cases:
713728
// - ResponseBodyLimit <= 0, which is the default behavior.
714-
// - [Request.SetOutputFile] is called to save response data to the file.
729+
// - [Request.SetOutputFileName] is called to save response data to the file.
715730
// - "DoNotParseResponse" is set for client or request.
716731
//
717732
// It overrides the value set at the client instance level, see [Client.SetResponseBodyLimit]

request_test.go

+49-5
Original file line numberDiff line numberDiff line change
@@ -1299,11 +1299,10 @@ func TestOutputFileWithBaseDirAndRelativePath(t *testing.T) {
12991299
SetRedirectPolicy(FlexibleRedirectPolicy(10)).
13001300
SetOutputDirectory(baseOutputDir).
13011301
SetDebug(true)
1302-
client.outputLogTo(io.Discard)
13031302

13041303
outputFilePath := "go-resty/test-img-success.png"
13051304
resp, err := client.R().
1306-
SetOutputFile(outputFilePath).
1305+
SetOutputFileName(outputFilePath).
13071306
Get(ts.URL + "/my-image.png")
13081307

13091308
assertError(t, err)
@@ -1332,7 +1331,7 @@ func TestOutputPathDirNotExists(t *testing.T) {
13321331
SetOutputDirectory(filepath.Join(getTestDataPath(), "not-exists-dir"))
13331332

13341333
resp, err := client.R().
1335-
SetOutputFile("test-img-success.png").
1334+
SetOutputFileName("test-img-success.png").
13361335
Get(ts.URL + "/my-image.png")
13371336

13381337
assertError(t, err)
@@ -1348,7 +1347,7 @@ func TestOutputFileAbsPath(t *testing.T) {
13481347
outputFile := filepath.Join(getTestDataPath(), "go-resty", "test-img-success-2.png")
13491348

13501349
res, err := dcnlr().
1351-
SetOutputFile(outputFile).
1350+
SetOutputFileName(outputFile).
13521351
Get(ts.URL + "/my-image.png")
13531352

13541353
assertError(t, err)
@@ -1358,6 +1357,51 @@ func TestOutputFileAbsPath(t *testing.T) {
13581357
assertNil(t, err)
13591358
}
13601359

1360+
func TestRequestSaveResponse(t *testing.T) {
1361+
ts := createGetServer(t)
1362+
defer ts.Close()
1363+
defer cleanupFiles(filepath.Join(".testdata", "go-resty"))
1364+
1365+
c := dcnl().
1366+
SetSaveResponse(true).
1367+
SetOutputDirectory(filepath.Join(getTestDataPath(), "go-resty"))
1368+
1369+
assertEqual(t, true, c.IsSaveResponse())
1370+
1371+
t.Run("content-disposition save response request", func(t *testing.T) {
1372+
outputFile := filepath.Join(getTestDataPath(), "go-resty", "test-img-success-2.png")
1373+
c.SetSaveResponse(false)
1374+
assertEqual(t, false, c.IsSaveResponse())
1375+
1376+
res, err := c.R().
1377+
SetSaveResponse(true).
1378+
Get(ts.URL + "/my-image.png?content-disposition=true&filename=test-img-success-2.png")
1379+
1380+
assertError(t, err)
1381+
assertEqual(t, int64(2579468), res.Size())
1382+
1383+
_, err = os.Stat(outputFile)
1384+
assertNil(t, err)
1385+
})
1386+
1387+
t.Run("use filename from path", func(t *testing.T) {
1388+
outputFile := filepath.Join(getTestDataPath(), "go-resty", "my-image.png")
1389+
c.SetSaveResponse(false)
1390+
assertEqual(t, false, c.IsSaveResponse())
1391+
1392+
res, err := c.R().
1393+
SetSaveResponse(true).
1394+
Get(ts.URL + "/my-image.png")
1395+
1396+
assertError(t, err)
1397+
assertEqual(t, int64(2579468), res.Size())
1398+
1399+
_, err = os.Stat(outputFile)
1400+
assertNil(t, err)
1401+
})
1402+
1403+
}
1404+
13611405
func TestContextInternal(t *testing.T) {
13621406
ts := createGetServer(t)
13631407
defer ts.Close()
@@ -2175,7 +2219,7 @@ func TestRequestSetResultAndSetOutputFile(t *testing.T) {
21752219
SetBody(&credentials{Username: "testuser", Password: "testpass"}).
21762220
SetResponseBodyUnlimitedReads(true).
21772221
SetResult(&AuthSuccess{}).
2178-
SetOutputFile(outputFile).
2222+
SetOutputFileName(outputFile).
21792223
Post("/login")
21802224

21812225
assertError(t, err)

resty_test.go

+4
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,10 @@ func createGetServer(t *testing.T) *httptest.Server {
114114
fileBytes, _ := os.ReadFile(filepath.Join(getTestDataPath(), "test-img.png"))
115115
w.Header().Set("Content-Type", "image/png")
116116
w.Header().Set("Content-Length", strconv.Itoa(len(fileBytes)))
117+
if r.URL.Query().Get("content-disposition") == "true" {
118+
filename := r.URL.Query().Get("filename")
119+
w.Header().Set(hdrContentDisposition, "inline; filename=\""+filename+"\"")
120+
}
117121
_, _ = w.Write(fileBytes)
118122
case "/get-method-payload-test":
119123
body, err := io.ReadAll(r.Body)

util.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -407,7 +407,7 @@ func responseDebugLogger(c *Client, res *Response) {
407407
fmt.Sprintf("TIME DURATION: %v\n", res.Time()) +
408408
"HEADERS :\n" +
409409
composeHeaders(rl.Header) + "\n"
410-
if res.Request.isSaveResponse {
410+
if res.Request.IsSaveResponse {
411411
debugLog += "BODY :\n***** RESPONSE WRITTEN INTO FILE *****\n"
412412
} else {
413413
debugLog += fmt.Sprintf("BODY :\n%v\n", rl.Body)

0 commit comments

Comments
 (0)