Repository: valyala/fasthttp
Pull request #631 (2019-08-12 20:50:16)

Add ability to set timeout for handshake
Relates to https://github.com/valyala/fasthttp/issues/630

mikefaraponov: Unfortunately, looks like SetWriteDeadline will not result in connection closed after timeout if server accepts connection and not closing it
mikefaraponov: Probably, another way to fix it SetReadDeadline earlier than SetWriteDeadline
erikdubbelboer: Looks like you are right. I mean more a test with code like this:
```go
package main

import (
	"net"
	"time"

	"github.com/valyala/fasthttp"
)

func main() {
	ln, err := net.Listen("tcp4", "127.0.0.1:1337")
	if err != nil {
		panic(err)
	}
	defer ln.Close()

	go func() {
		for {
			_, err := ln.Accept()
			if err != nil {
				panic(err)
			}

			println("connection accepted")
		}
	}()

	req := fasthttp.AcquireRequest()
	res := fasthttp.AcquireResponse()

	req.SetRequestURIBytes([]byte("https://127.0.0.1:1337/"))

	client := fasthttp.Client{
		WriteTimeout: time.Second,
		ReadTimeout:  time.Second,
	}

	if err := client.Do(req, res); err != nil {
		panic(err)
	}

	println("this should happen in 1 second, but it doesn't")
}
```
If you can turn this into a test case that will fail before your pull request but properly timeout after your pull request that would be great.

I think I'm in favor of setting the TLS timeout automatically to the same value as `WriteTimeout`. Since this time is included in trying to write the request.
mikefaraponov: I think it's a good idea.

Patch:
diff --git a/client.go b/client.go
index cd55961..9ca53d2 100644
--- a/client.go
+++ b/client.go
@@ -1537,7 +1537,7 @@ func (c *HostClient) dialHostHard() (conn net.Conn, err error) {
 	for n > 0 {
 		addr := c.nextAddr()
 		tlsConfig := c.cachedTLSConfig(addr)
-		conn, err = dialAddr(addr, c.Dial, c.DialDualStack, c.IsTLS, tlsConfig)
+		conn, err = dialAddr(addr, c.Dial, c.DialDualStack, c.IsTLS, tlsConfig, c.WriteTimeout)
 		if err == nil {
 			return conn, nil
 		}
@@ -1568,7 +1568,43 @@ func (c *HostClient) cachedTLSConfig(addr string) *tls.Config {
 	return cfg
 }
 
-func dialAddr(addr string, dial DialFunc, dialDualStack, isTLS bool, tlsConfig *tls.Config) (net.Conn, error) {
+var ErrTLSHandshakeTimeout = errors.New("tls handshake timed out")
+
+var timeoutErrorChPool sync.Pool
+
+func tlsClientHandshake(rawConn net.Conn, tlsConfig *tls.Config, timeout time.Duration) (net.Conn, error) {
+	tc := AcquireTimer(timeout)
+	defer ReleaseTimer(tc)
+
+	var ch chan error
+	chv := timeoutErrorChPool.Get()
+	if chv == nil {
+		chv = make(chan error)
+	}
+	ch = chv.(chan error)
+	defer timeoutErrorChPool.Put(chv)
+
+	conn := tls.Client(rawConn, tlsConfig)
+
+	go func() {
+		ch <- conn.Handshake()
+	}()
+
+	select {
+	case <-tc.C:
+		rawConn.Close()
+		<-ch
+		return nil, ErrTLSHandshakeTimeout
+	case err := <-ch:
+		if err != nil {
+			rawConn.Close()
+			return nil, err
+		}
+		return conn, nil
+	}
+}
+
+func dialAddr(addr string, dial DialFunc, dialDualStack, isTLS bool, tlsConfig *tls.Config, timeout time.Duration) (net.Conn, error) {
 	if dial == nil {
 		if dialDualStack {
 			dial = DialDualStack
@@ -1585,7 +1621,10 @@ func dialAddr(addr string, dial DialFunc, dialDualStack, isTLS bool, tlsConfig *
 		panic("BUG: DialFunc returned (nil, nil)")
 	}
 	if isTLS {
-		conn = tls.Client(conn, tlsConfig)
+		if timeout == 0 {
+			return tls.Client(conn, tlsConfig), nil
+		}
+		return tlsClientHandshake(conn, tlsConfig, timeout)
 	}
 	return conn, nil
 }
@@ -1992,7 +2031,7 @@ func (c *pipelineConnClient) init() {
 
 func (c *pipelineConnClient) worker() error {
 	tlsConfig := c.cachedTLSConfig()
-	conn, err := dialAddr(c.Addr, c.Dial, c.DialDualStack, c.IsTLS, tlsConfig)
+	conn, err := dialAddr(c.Addr, c.Dial, c.DialDualStack, c.IsTLS, tlsConfig, c.WriteTimeout)
 	if err != nil {
 		return err
 	}
diff --git a/http.go b/http.go
index 737673a..99782a8 100644
--- a/http.go
+++ b/http.go
@@ -3,6 +3,7 @@ package fasthttp
 import (
 	"bufio"
 	"bytes"
+	"encoding/base64"
 	"errors"
 	"fmt"
 	"io"
@@ -1148,6 +1149,24 @@ func (req *Request) Write(w *bufio.Writer) error {
 		}
 		req.Header.SetHostBytes(host)
 		req.Header.SetRequestURIBytes(uri.RequestURI())
+
+		if len(uri.username) > 0 {
+			// RequestHeader.SetBytesKV only uses RequestHeader.bufKV.key
+			// So we are free to use RequestHeader.bufKV.value as a scratch pad for
+			// the base64 encoding.
+			nl := len(uri.username) + len(uri.password) + 1
+			tl := nl + base64.StdEncoding.EncodedLen(nl)
+			if tl > cap(req.Header.bufKV.value) {
+				req.Header.bufKV.value = make([]byte, 0, tl)
+			}
+			buf := req.Header.bufKV.value[:0]
+			buf = append(buf, uri.username...)
+			buf = append(buf, strColon...)
+			buf = append(buf, uri.password...)
+			buf = buf[:tl]
+			base64.StdEncoding.Encode(buf[nl:], buf[:nl])
+			req.Header.SetBytesKV(strAuthorization, buf[nl:])
+		}
 	}
 
 	if req.bodyStream != nil {
diff --git a/strings.go b/strings.go
index f654f95..343544a 100644
--- a/strings.go
+++ b/strings.go
@@ -16,9 +16,11 @@ var (
 	strHTTP             = []byte("http")
 	strHTTPS            = []byte("https")
 	strHTTP11           = []byte("HTTP/1.1")
+	strColon            = []byte(":")
 	strColonSlashSlash  = []byte("://")
 	strColonSpace       = []byte(": ")
 	strGMT              = []byte("GMT")
+	strAt               = []byte("@")
 
 	strResponseContinue = []byte("HTTP/1.1 100 Continue\r\n\r\n")
 
@@ -52,6 +54,7 @@ var (
 	strAcceptRanges     = []byte(HeaderAcceptRanges)
 	strRange            = []byte(HeaderRange)
 	strContentRange     = []byte(HeaderContentRange)
+	strAuthorization    = []byte(HeaderAuthorization)
 
 	strCookieExpires        = []byte("expires")
 	strCookieDomain         = []byte("domain")
diff --git a/uri.go b/uri.go
index d536f59..e9cd4b1 100644
--- a/uri.go
+++ b/uri.go
@@ -51,6 +51,9 @@ type URI struct {
 	fullURI    []byte
 	requestURI []byte
 
+	username []byte
+	password []byte
+
 	h *RequestHeader
 }
 
@@ -63,6 +66,8 @@ func (u *URI) CopyTo(dst *URI) {
 	dst.queryString = append(dst.queryString[:0], u.queryString...)
 	dst.hash = append(dst.hash[:0], u.hash...)
 	dst.host = append(dst.host[:0], u.host...)
+	dst.username = append(dst.username[:0], u.username...)
+	dst.password = append(dst.password[:0], u.password...)
 
 	u.queryArgs.CopyTo(&dst.queryArgs)
 	dst.parsedQueryArgs = u.parsedQueryArgs
@@ -89,6 +94,36 @@ func (u *URI) SetHashBytes(hash []byte) {
 	u.hash = append(u.hash[:0], hash...)
 }
 
+// Username returns URI username
+func (u *URI) Username() []byte {
+	return u.username
+}
+
+// SetUsername sets URI username.
+func (u *URI) SetUsername(username string) {
+	u.username = append(u.username[:0], username...)
+}
+
+// SetUsernameBytes sets URI username.
+func (u *URI) SetUsernameBytes(username []byte) {
+	u.username = append(u.username[:0], username...)
+}
+
+// Password returns URI password
+func (u *URI) Password() []byte {
+	return u.password
+}
+
+// SetPassword sets URI password.
+func (u *URI) SetPassword(password string) {
+	u.password = append(u.password[:0], password...)
+}
+
+// SetPasswordBytes sets URI password.
+func (u *URI) SetPasswordBytes(password []byte) {
+	u.password = append(u.password[:0], password...)
+}
+
 // QueryString returns URI query string,
 // i.e. baz=123 of http://aaa.com/foo/bar?baz=123#qwe .
 //
@@ -174,6 +209,8 @@ func (u *URI) Reset() {
 	u.path = u.path[:0]
 	u.queryString = u.queryString[:0]
 	u.hash = u.hash[:0]
+	u.username = u.username[:0]
+	u.password = u.password[:0]
 
 	u.host = u.host[:0]
 	u.queryArgs.Reset()
@@ -236,6 +273,20 @@ func (u *URI) parse(host, uri []byte, h *RequestHeader) {
 	scheme, host, uri := splitHostURI(host, uri)
 	u.scheme = append(u.scheme, scheme...)
 	lowercaseBytes(u.scheme)
+
+	if n := bytes.Index(host, strAt); n >= 0 {
+		auth := host[:n]
+		host = host[n+1:]
+
+		if n := bytes.Index(auth, strColon); n >= 0 {
+			u.username = auth[:n]
+			u.password = auth[n+1:]
+		} else {
+			u.username = auth
+			u.password = auth[:0] // Make sure it's not nil
+		}
+	}
+
 	u.host = append(u.host, host...)
 	lowercaseBytes(u.host)
 


Test patch:
diff --git a/allocation_test.go b/allocation_test.go
index e57c3af..45f10e5 100644
--- a/allocation_test.go
+++ b/allocation_test.go
@@ -49,7 +49,7 @@ func TestAllocationClient(t *testing.T) {
 	go s.Serve(ln)
 
 	c := &Client{}
-	url := "http://" + ln.Addr().String()
+	url := "http://test:test@" + ln.Addr().String() + "/foo?bar=baz"
 
 	n := testing.AllocsPerRun(100, func() {
 		req := AcquireRequest()
@@ -68,3 +68,17 @@ func TestAllocationClient(t *testing.T) {
 		t.Fatalf("expected 0 allocations, got %f", n)
 	}
 }
+
+func TestAllocationURI(t *testing.T) {
+	uri := []byte("http://username:password@example.com/some/path?foo=bar#test")
+
+	n := testing.AllocsPerRun(100, func() {
+		u := AcquireURI()
+		u.Parse(nil, uri)
+		ReleaseURI(u)
+	})
+
+	if n != 0 {
+		t.Fatalf("expected 0 allocations, got %f", n)
+	}
+}
diff --git a/client_test.go b/client_test.go
index d602b7d..9871350 100644
--- a/client_test.go
+++ b/client_test.go
@@ -19,6 +19,44 @@ import (
 	"github.com/valyala/fasthttp/fasthttputil"
 )
 
+func TestClientURLAuth(t *testing.T) {
+	cases := map[string]string{
+		"user:pass@": "dXNlcjpwYXNz",
+		"foo:@":      "Zm9vOg==",
+		":@":         "",
+		"@":          "",
+		"":           "",
+	}
+
+	ch := make(chan string, 1)
+	ln := fasthttputil.NewInmemoryListener()
+	s := &Server{
+		Handler: func(ctx *RequestCtx) {
+			ch <- string(ctx.Request.Header.Peek(HeaderAuthorization))
+		},
+	}
+	go s.Serve(ln)
+	c := &Client{
+		Dial: func(addr string) (net.Conn, error) {
+			return ln.Dial()
+		},
+	}
+	for up, expected := range cases {
+		req := AcquireRequest()
+		req.Header.SetMethod(MethodGet)
+		req.SetRequestURI("http://" + up + "example.com")
+		if err := c.Do(req, nil); err != nil {
+			t.Fatal(err)
+		}
+
+		val := <-ch
+
+		if val != expected {
+			t.Fatalf("wrong %s header: %s expected %s", HeaderAuthorization, val, expected)
+		}
+	}
+}
+
 func TestClientNilResp(t *testing.T) {
 	ln := fasthttputil.NewInmemoryListener()
 	s := &Server{
@@ -1735,3 +1773,43 @@ func startEchoServerExt(t *testing.T, network, addr string, isTLS bool) *testEch
 		t:  t,
 	}
 }
+
+func TestClientTLSHandshakeTimeout(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping test in short mode")
+	}
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	addr := listener.Addr().String()
+	defer listener.Close()
+
+	complete := make(chan bool)
+	defer close(complete)
+
+	go func() {
+		conn, err := listener.Accept()
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		<-complete
+		conn.Close()
+	}()
+
+	client := Client{
+		WriteTimeout: 1 * time.Second,
+		ReadTimeout:  1 * time.Second,
+	}
+
+	_, _, err = client.Get(nil, "https://"+addr)
+	if err == nil {
+		t.Fatal("tlsClientHandshake completed successfully")
+	}
+
+	if err != ErrTLSHandshakeTimeout {
+		t.Errorf("resulting error not a timeout: %v\nType %T: %#v", err, err, err)
+	}
+}