Repository: valyala/fasthttp
Pull request #1415 (2022-11-02 21:56:01)

Request.SetTimeout
This functionally works the same as e.g. Client.DoTimeout(), but we can also use it for Client.DoRedirect(). There is no way as far as I can tell to set a timeout on a DoRedirect call, so this makes it possible.

li-jin-gou: Can you add some tests to test the behavior you would expect from this function and some comment? 
brian-armstrong-discord: @li-jin-gou Thanks for the feedback! I've added some tests both for the feature in isolation (in `client.Do()`) and in the context of Redirects. It looks to be working as expected.

Patch:
diff --git a/http.go b/http.go
index 62ba30f..4b9e113 100644
--- a/http.go
+++ b/http.go
@@ -2001,7 +2001,7 @@ func writeBodyChunked(w *bufio.Writer, r io.Reader) error {
 		n, err = r.Read(buf)
 		if n == 0 {
 			if err == nil {
-				panic("BUG: io.Reader returned 0, nil")
+				continue
 			}
 			if err == io.EOF {
 				if err = writeChunk(w, buf[:0]); err != nil {
@@ -2290,3 +2290,11 @@ func round2(n int) int {
 
 	return int(x + 1)
 }
+
+// SetTimeout sets timeout for the request.
+//
+// req.SetTimeout(t); c.Do(&req, &resp) is equivalent to
+// c.DoTimeout(&req, &resp, t)
+func (req *Request) SetTimeout(t time.Duration) {
+	req.timeout = t
+}


Test patch:
diff --git a/client_test.go b/client_test.go
index 90092f4..13da521 100644
--- a/client_test.go
+++ b/client_test.go
@@ -1566,6 +1566,56 @@ func TestClientFollowRedirects(t *testing.T) {
 		ReleaseResponse(resp)
 	}
 
+	for i := 0; i < 10; i++ {
+		req := AcquireRequest()
+		resp := AcquireResponse()
+
+		req.SetRequestURI("http://xxx/foo")
+
+		req.SetTimeout(time.Second)
+		err := c.DoRedirects(req, resp, 16)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+
+		if statusCode := resp.StatusCode(); statusCode != StatusOK {
+			t.Fatalf("unexpected status code: %d", statusCode)
+		}
+
+		if body := string(resp.Body()); body != "/bar" {
+			t.Fatalf("unexpected response %q. Expecting %q", body, "/bar")
+		}
+
+		ReleaseRequest(req)
+		ReleaseResponse(resp)
+	}
+
+	for i := 0; i < 10; i++ {
+		req := AcquireRequest()
+		resp := AcquireResponse()
+
+		req.SetRequestURI("http://xxx/foo")
+
+		testConn, _ := net.Dial("tcp", ln.Addr().String())
+		timeoutConn := &Client{
+			Dial: func(addr string) (net.Conn, error) {
+				return &readTimeoutConn{Conn: testConn, t: time.Second}, nil
+			},
+		}
+
+		req.SetTimeout(time.Millisecond)
+		err := timeoutConn.DoRedirects(req, resp, 16)
+		if err == nil {
+			t.Errorf("expecting error")
+		}
+		if err != ErrTimeout {
+			t.Errorf("unexpected error: %v. Expecting %v", err, ErrTimeout)
+		}
+
+		ReleaseRequest(req)
+		ReleaseResponse(resp)
+	}
+
 	req := AcquireRequest()
 	resp := AcquireResponse()
 
@@ -1613,6 +1663,7 @@ func TestClientDoTimeoutSuccess(t *testing.T) {
 	defer s.Stop()
 
 	testClientDoTimeoutSuccess(t, &defaultClient, "http://"+s.Addr(), 100)
+	testClientRequestSetTimeoutSuccess(t, &defaultClient, "http://"+s.Addr(), 100)
 }
 
 func TestClientDoTimeoutSuccessConcurrent(t *testing.T) {
@@ -1627,6 +1678,7 @@ func TestClientDoTimeoutSuccessConcurrent(t *testing.T) {
 		go func() {
 			defer wg.Done()
 			testClientDoTimeoutSuccess(t, &defaultClient, "http://"+s.Addr(), 100)
+			testClientRequestSetTimeoutSuccess(t, &defaultClient, "http://"+s.Addr(), 100)
 		}()
 	}
 	wg.Wait()
@@ -1687,6 +1739,7 @@ func TestClientDoTimeoutError(t *testing.T) {
 	}
 
 	testClientDoTimeoutError(t, c, 100)
+	testClientRequestSetTimeoutError(t, c, 100)
 }
 
 func TestClientDoTimeoutErrorConcurrent(t *testing.T) {
@@ -1748,6 +1801,22 @@ func testClientGetTimeoutError(t *testing.T, c *Client, n int) {
 	}
 }
 
+func testClientRequestSetTimeoutError(t *testing.T, c *Client, n int) {
+	var req Request
+	var resp Response
+	req.SetRequestURI("http://foobar.com/baz")
+	for i := 0; i < n; i++ {
+		req.SetTimeout(time.Millisecond)
+		err := c.Do(&req, &resp)
+		if err == nil {
+			t.Errorf("expecting error")
+		}
+		if err != ErrTimeout {
+			t.Errorf("unexpected error: %v. Expecting %v", err, ErrTimeout)
+		}
+	}
+}
+
 type readTimeoutConn struct {
 	net.Conn
 	t  time.Duration
@@ -2398,6 +2467,30 @@ func testClientDoTimeoutSuccess(t *testing.T, c *Client, addr string, n int) {
 	}
 }
 
+func testClientRequestSetTimeoutSuccess(t *testing.T, c *Client, addr string, n int) {
+	var req Request
+	var resp Response
+
+	for i := 0; i < n; i++ {
+		uri := fmt.Sprintf("%s/foo/%d?bar=baz", addr, i)
+		req.SetRequestURI(uri)
+		req.SetTimeout(time.Second)
+		if err := c.Do(&req, &resp); err != nil {
+			t.Errorf("unexpected error: %v", err)
+		}
+		if resp.StatusCode() != StatusOK {
+			t.Errorf("unexpected status code: %d. Expecting %d", resp.StatusCode(), StatusOK)
+		}
+		resultURI := string(resp.Body())
+		if strings.HasPrefix(uri, "https") {
+			resultURI = uri[:5] + resultURI[4:]
+		}
+		if resultURI != uri {
+			t.Errorf("unexpected uri %q. Expecting %q", resultURI, uri)
+		}
+	}
+}
+
 func testClientGetTimeoutSuccess(t *testing.T, c *Client, addr string, n int) {
 	var buf []byte
 	for i := 0; i < n; i++ {
diff --git a/http_test.go b/http_test.go
index 60e7891..52d8a67 100644
--- a/http_test.go
+++ b/http_test.go
@@ -2949,3 +2949,87 @@ func TestResponseBodyStreamErrorOnPanicDuringClose(t *testing.T) {
 		t.Fatalf("unexpected error value, got: %+v.", e.Error())
 	}
 }
+
+func TestRequestMultipartFormPipeEmptyFormField(t *testing.T) {
+	t.Parallel()
+
+	pr, pw := io.Pipe()
+	mw := multipart.NewWriter(pw)
+
+	errs := make(chan error, 1)
+
+	go func() {
+		defer func() {
+			err := mw.Close()
+			if err != nil {
+				errs <- err
+			}
+			err = pw.Close()
+			if err != nil {
+				errs <- err
+			}
+			close(errs)
+		}()
+
+		if err := mw.WriteField("emptyField", ""); err != nil {
+			errs <- err
+		}
+	}()
+
+	var b bytes.Buffer
+	bw := bufio.NewWriter(&b)
+	err := writeBodyChunked(bw, pr)
+
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	for e := range errs {
+		t.Fatalf("unexpected error in goroutine multiwriter: %v", e)
+	}
+
+	testRequestMultipartFormPipeEmptyFormField(t, mw.Boundary(), b.Bytes(), 1)
+}
+
+func testRequestMultipartFormPipeEmptyFormField(t *testing.T, boundary string, formData []byte, partsCount int) []byte {
+	s := fmt.Sprintf("POST / HTTP/1.1\r\nHost: aaa\r\nContent-Type: multipart/form-data; boundary=%s\r\nTransfer-Encoding: chunked\r\n\r\n%s",
+		boundary, formData)
+
+	var req Request
+
+	r := bytes.NewBufferString(s)
+	br := bufio.NewReader(r)
+	if err := req.Read(br); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	f, err := req.MultipartForm()
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	defer req.RemoveMultipartFormFiles()
+
+	if len(f.File) > 0 {
+		t.Fatalf("unexpected files found in the multipart form: %d", len(f.File))
+	}
+
+	if len(f.Value) != partsCount {
+		t.Fatalf("unexpected number of values found: %d. Expecting %d", len(f.Value), partsCount)
+	}
+
+	for k, vv := range f.Value {
+		if len(vv) != 1 {
+			t.Fatalf("unexpected number of values found for key=%q: %d. Expecting 1", k, len(vv))
+		}
+		if k != "emptyField" {
+			t.Fatalf("unexpected key=%q. Expecting %q", k, "emptyField")
+		}
+
+		v := vv[0]
+		if v != "" {
+			t.Fatalf("unexpected value=%q. expecting %q", v, "")
+		}
+	}
+
+	return req.Body()
+}