Supports saving large file size too, code & test improvements.

Author: jeevatkm

client.go

@@ -571,10 +571,14 @@ func (c *Client) execute(req *Request) (*Response, error) {
 		RawResponse: resp,
 	}
 
-	defer resp.Body.Close()
-	response.Body, err = ioutil.ReadAll(resp.Body)
-	if err != nil {
-		return nil, err
+	if !req.isSaveResponse {
+		defer resp.Body.Close()
+		response.Body, err = ioutil.ReadAll(resp.Body)
+		if err != nil {
+			return nil, err
+		}
+
+		response.size = int64(len(response.Body))
 	}
 
 	// Apply Response middleware
@@ -890,6 +894,7 @@ func (r *Request) SetAuthToken(token string) *Request {
 // 			SetOutput("/Users/jeeva/Downloads/ReplyWithHeader-v5.1-beta.zip").
 // 			Get("http://bit.ly/1LouEKr")
 //
+// Note: In this scenario `Response.Body` might be nil.
 func (r *Request) SetOutput(file string) *Request {
 	r.outputFile = file
 	r.isSaveResponse = true
@@ -960,6 +965,8 @@ type Response struct {
 	ReceivedAt  time.Time
 	Request     *Request
 	RawResponse *http.Response
+
+	size int64
 }
 
 // Status method returns the HTTP status string for the executed request.
@@ -1010,6 +1017,13 @@ func (r *Response) Time() time.Duration {
 	return r.ReceivedAt.Sub(r.Request.Time)
 }
 
+// Size method returns the HTTP response size in bytes. Ya, you can relay on HTTP `Content-Length` header,
+// however it won't be good for chucked transfer/compressed response. Since Resty calculates response size
+// at the client end. You will get actual size of the http response.
+func (r *Response) Size() int64 {
+	return r.size
+}
+
 //
 // Helper methods
 //
@@ -1167,7 +1181,7 @@ func createDirectory(dir string) (err error) {
 	if _, err = os.Stat(dir); err != nil {
 		if os.IsNotExist(err) {
 			if err = os.MkdirAll(dir, 0755); err != nil {
-				return err
+				return
 			}
 		}
 	}

middleware.go

@@ -10,10 +10,11 @@ import (
 	"encoding/xml"
 	"errors"
 	"fmt"
-	"io/ioutil"
+	"io"
 	"mime/multipart"
 	"net/http"
 	"net/url"
+	"os"
 	"path/filepath"
 	"reflect"
 	"strings"
@@ -250,7 +251,11 @@ func responseLogger(c *Client, res *Response) error {
 		for h, v := range res.Header() {
 			c.Log.Printf("%30s: %v", h, strings.Join(v, ", "))
 		}
-		c.Log.Printf("BODY   :\n%v", getResponseBodyString(res))
+		if res.Request.isSaveResponse {
+			c.Log.Printf("BODY   :\n***** RESPONSE WRITTEN INTO FILE *****")
+		} else {
+			c.Log.Printf("BODY   :\n%v", getResponseBodyString(res))
+		}
 		c.Log.Println("----------------------------------------------------------")
 		c.enableLogPrefix()
 	}
@@ -285,27 +290,35 @@ func parseResponseBody(c *Client, res *Response) (err error) {
 	return
 }
 
-func saveResponseIntoFile(c *Client, res *Response) (err error) {
+func saveResponseIntoFile(c *Client, res *Response) error {
 	if res.Request.isSaveResponse {
 		file := ""
 
 		if len(c.outputDirectory) > 0 && !filepath.IsAbs(res.Request.outputFile) {
 			file += c.outputDirectory + string(filepath.Separator)
 		}
+
 		file = filepath.Clean(file + res.Request.outputFile)
+		err := createDirectory(filepath.Dir(file))
+		if err != nil {
+			return err
+		}
 
-		err = createDirectory(filepath.Dir(file))
+		outFile, err := os.Create(file)
 		if err != nil {
-			c.Log.Printf("ERROR [%v]", err)
-			return
+			return err
 		}
+		defer outFile.Close()
 
-		err = ioutil.WriteFile(file, res.Body, 0644)
+		// io.Copy reads maximum 32kb size, it is perfect for large file download too
+		defer res.RawResponse.Body.Close()
+		written, err := io.Copy(outFile, res.RawResponse.Body)
 		if err != nil {
-			c.Log.Printf("ERROR [%v]", err)
-			return
+			return err
 		}
+
+		res.size = written
 	}
 
-	return
+	return nil
 }

redirect.go

@@ -50,6 +50,7 @@ func FlexibleRedirectPolicy(noOfRedirect int) RedirectPolicy {
 }
 
 // DomainCheckRedirectPolicy is convenient method to define domain name redirect rule in resty client.
+// Redirect is allowed for only mentioned host in the policy.
 // 		resty.SetRedirectPolicy(DomainCheckRedirectPolicy("host1.com", "host2.org", "host3.net"))
 func DomainCheckRedirectPolicy(hostnames ...string) RedirectPolicy {
 	hosts := make(map[string]bool)
@@ -58,9 +59,16 @@ func DomainCheckRedirectPolicy(hostnames ...string) RedirectPolicy {
 	}
 
 	fn := RedirectPolicyFunc(func(req *http.Request, via []*http.Request) error {
-		host, _, _ := net.SplitHostPort(req.URL.Host)
-		if _, ok := hosts[strings.ToLower(host)]; ok {
-			return fmt.Errorf("Redirect is not allowed for host[%s]", strings.ToLower(host))
+		hostname := ""
+		if strings.Index(req.URL.Host, ":") > 0 {
+			host, _, _ := net.SplitHostPort(req.URL.Host)
+			hostname = strings.ToLower(host)
+		} else {
+			hostname = strings.ToLower(req.URL.Host)
+		}
+
+		if ok := hosts[hostname]; !ok {
+			return errors.New("Redirect is not allowed as per DomainCheckRedirectPolicy")
 		}
 
 		return nil

resty.go

@@ -9,4 +9,4 @@ Package resty provides simple HTTP and REST client for Go inspired by Ruby rest-
 package resty
 
 // go-resty version no
-var Version = "0.4"
+var Version = "0.4.1"

resty_test.go

@@ -733,10 +733,12 @@ func TestHostCheckRedirectPolicy(t *testing.T) {
 	defer ts.Close()
 
 	c := dc().
-		SetRedirectPolicy(DomainCheckRedirectPolicy("NotAllowed.com")).
-		SetTimeout(time.Duration(time.Second * 2))
+		SetRedirectPolicy(DomainCheckRedirectPolicy("127.0.0.1"))
 
-	c.R().Get(ts.URL + "/redirect-host-check-1")
+	_, err := c.R().Get(ts.URL + "/redirect-host-check-1")
+
+	assertEqual(t, true, err != nil)
+	assertEqual(t, true, strings.Contains(err.Error(), "Redirect is not allowed as per DomainCheckRedirectPolicy"))
 }
 
 func TestClientRedirectPolicy(t *testing.T) {
@@ -917,54 +919,39 @@ func TestSetRootCertificateNotExists(t *testing.T) {
 	assertEqual(t, true, DefaultClient.transport.TLSClientConfig == nil)
 }
 
-func TestOutputFileRelativePath(t *testing.T) {
-	c := dc().SetRedirectPolicy(FlexibleRedirectPolicy(10))
-
-	_, err := c.R().
-		SetOutput("go-resty/ReplyWithHeader-v5.1-beta.zip").
-		Get("http://bit.ly/1LouEKr")
-
-	assertError(t, err)
-}
+func TestOutputFileWithBaseDirAndRelativePath(t *testing.T) {
+	ts := createGetServer(t)
+	defer ts.Close()
 
-func TestOutputFileWithBaseDir(t *testing.T) {
 	DefaultClient = dc()
-
 	SetRedirectPolicy(FlexibleRedirectPolicy(10))
-	SetOutputDirectory(getTestDataPath() + "/sample1")
-
-	_, err := R().
-		SetOutput("go-resty/ReplyWithHeader-v5.1-beta.zip").
-		Get("http://bit.ly/1LouEKr")
-
-	assertError(t, err)
-}
-
-func TestOutputFileAbsPath(t *testing.T) {
-	c := dc().SetRedirectPolicy(FlexibleRedirectPolicy(10))
+	SetOutputDirectory(getTestDataPath() + "/dir-sample")
+	SetDebug(true)
 
-	_, err := c.R().
-		SetOutput(getTestDataPath() + "/go-resty-test-2/ReplyWithHeader-v5.1-beta.zip").
-		Get("http://bit.ly/1LouEKr")
+	resp, err := R().
+		SetOutput("go-resty/test-img-success.png").
+		Get(ts.URL + "/my-image.png")
 
 	assertError(t, err)
+	assertEqual(t, true, resp.Size() != 0)
 }
 
 func TestOutputFileWithBaseDirError(t *testing.T) {
 	c := dc().SetRedirectPolicy(FlexibleRedirectPolicy(10)).
-		SetOutputDirectory("/go-resty")
+		SetOutputDirectory(getTestDataPath() + `/go-resty\0`)
 
 	_ = c
 }
 
-func TestOutputFileAbsPathError(t *testing.T) {
-	c := dc().SetRedirectPolicy(FlexibleRedirectPolicy(10))
+func TestOutputFileAbsPath(t *testing.T) {
+	ts := createGetServer(t)
+	defer ts.Close()
 
-	_, err := c.R().
-		SetOutput("/go-resty-test-2/ReplyWithHeader-v5.1-beta.zip").
-		Get("http://bit.ly/1LouEKr")
+	_, err := dcr().
+		SetOutput(getTestDataPath() + "/go-resty/test-img-success-2.png").
+		Get(ts.URL + "/my-image.png")
 
-	assertEqual(t, true, err != nil)
+	assertError(t, err)
 }
 
 func TestClientOptions(t *testing.T) {
@@ -1071,6 +1058,11 @@ func createGetServer(t *testing.T) *httptest.Server {
 			} else if r.URL.Path == "/set-timeout-test" {
 				time.Sleep(time.Second * 5)
 				w.Write([]byte("TestClientTimeout page"))
+			} else if r.URL.Path == "/my-image.png" {
+				fileBytes, _ := ioutil.ReadFile(getTestDataPath() + "/test-img.png")
+				w.Header().Set("Content-Type", "image/png")
+				w.Header().Set("Content-Length", strconv.Itoa(len(fileBytes)))
+				w.Write(fileBytes)
 			}
 		}
 	})
@@ -1301,7 +1293,7 @@ func createRedirectServer(t *testing.T) *httptest.Server {
 
 				if cnt != 7 { // Testing hard stop via logical
 					if cnt >= 5 {
-						http.Redirect(w, r, "http://notallowed.com/go-redirect", http.StatusTemporaryRedirect)
+						http.Redirect(w, r, "http://httpbin.org/get", http.StatusTemporaryRedirect)
 					} else {
 						http.Redirect(w, r, fmt.Sprintf("/redirect-host-check-%d", (cnt+1)), http.StatusTemporaryRedirect)
 					}