修改1.timeout，2. 修改测试代码 (#28)

Author: guonaihong

client.go

@@ -196,34 +196,32 @@ func (d *DialOption) Dial() (wsCon *Conn, err error) {
 	}
 
 	var conn net.Conn
-	begin := time.Now()
 
 	hostName := hostname.GetHostName(d.u)
-	// conn, err := net.DialTimeout("tcp", d.u.Host /* TODO 加端号*/, d.dialTimeout)
-	dialFunc := net.Dial
+	dialFunc := net.DialTimeout
 	if d.dialFunc != nil {
 		dialInterface, err := d.dialFunc()
 		if err != nil {
 			return nil, err
 		}
-		dialFunc = dialInterface.Dial
+		dialFunc = func(network, address string, timeout time.Duration) (net.Conn, error) {
+			return dialInterface.Dial(network, address)
+		}
 	}
 
 	if d.proxyFunc != nil {
 		proxyURL, err := d.proxyFunc(req)
 		if err != nil {
 			return nil, err
 		}
-		dialFunc = newhttpProxy(proxyURL, dialFunc).Dial
+		dialFunc = newhttpProxy(proxyURL, dialFunc).DialTimeout
 	}
 
-	conn, err = dialFunc("tcp", hostName)
+	conn, err = dialFunc("tcp", hostName, d.dialTimeout)
 	if err != nil {
 		return nil, err
 	}
 
-	dialDuration := time.Since(begin)
-
 	conn = d.tlsConn(conn)
 	defer func() {
 		if err != nil && conn != nil {
@@ -232,18 +230,7 @@ func (d *DialOption) Dial() (wsCon *Conn, err error) {
 		}
 	}()
 
-	if to := d.dialTimeout - dialDuration; to > 0 {
-		if err = conn.SetDeadline(time.Now().Add(to)); err != nil {
-			return
-		}
-	}
-
-	defer func() {
-		if err == nil {
-			err = conn.SetDeadline(time.Time{})
-		}
-	}()
-
+	err = conn.SetDeadline(time.Time{})
 	if err = req.Write(conn); err != nil {
 		return
 	}

common_options_test.go

@@ -2227,13 +2227,14 @@ func Test_CommonOption(t *testing.T) {
 	t.Run("22.3.WithClientReadMaxMessage", func(t *testing.T) {
 		var tsort testServerOptionReadTimeout
 
-		upgrade := NewUpgrade(WithServerCallback(&tsort), WithServerReadTimeout(time.Millisecond*60))
+		upgrade := NewUpgrade()
 		tsort.err = make(chan error, 1)
 		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			c, err := upgrade.Upgrade(w, r)
 			if err != nil {
 				t.Error(err)
 			}
+			time.Sleep(time.Second / 100)
 			err = c.WriteMessage(Binary, bytes.Repeat([]byte("1"), 1025))
 			if err != nil {
 				t.Error(err)
@@ -2245,12 +2246,57 @@ func Test_CommonOption(t *testing.T) {
 		defer ts.Close()
 
 		url := strings.ReplaceAll(ts.URL, "http", "ws")
-		con, err := Dial(url, WithClientBufioParseMode(), WithClientReadMaxMessage(1<<10), WithClientOnMessageFunc(func(c *Conn, mt Opcode, payload []byte) {
+		con, err := Dial(url, WithClientCallback(&tsort), WithClientBufioParseMode(), WithClientReadMaxMessage(1<<10))
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		defer con.Close()
+		go func() {
+			_ = con.ReadLoop()
+		}()
+
+		select {
+		case d := <-tsort.err:
+			if d == nil {
+				t.Errorf("got:nil, need:error\n")
+			}
+		case <-time.After(100 * time.Hour):
+			t.Errorf(" Test_ServerOption:WithServerReadMaxMessage timeout\n")
+		}
+		if atomic.LoadInt32(&tsort.run) != 1 {
+			t.Error("not run server:method fail")
+		}
+	})
+	t.Run("22.4.WithClientReadMaxMessage-ParseWindows", func(t *testing.T) {
+		var tsort testServerOptionReadTimeout
+
+		upgrade := NewUpgrade()
+		tsort.err = make(chan error, 1)
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			c, err := upgrade.Upgrade(w, r)
+			if err != nil {
+				t.Error(err)
+			}
+			time.Sleep(time.Second / 100)
+			err = c.WriteMessage(Binary, bytes.Repeat([]byte("1"), 1025))
+			if err != nil {
+				t.Error(err)
+				return
+			}
+			c.StartReadLoop()
 		}))
+
+		defer ts.Close()
+
+		url := strings.ReplaceAll(ts.URL, "http", "ws")
+		con, err := Dial(url, WithClientCallback(&tsort), WithClientReadMaxMessage(1<<10))
 		if err != nil {
 			t.Error(err)
+			return
 		}
 		defer con.Close()
+		con.StartReadLoop()
 
 		select {
 		case d := <-tsort.err:

config.go

@@ -27,10 +27,16 @@ import (
 
 var ErrDialFuncAndProxyFunc = errors.New("dialFunc and proxyFunc can't be set at the same time")
 
+// 握手
 type Dialer interface {
 	Dial(network, addr string) (c net.Conn, err error)
 }
 
+// 带超时时间的握手
+type DialerTimeout interface {
+	DialTimeout(network, addr string, timeout time.Duration) (c net.Conn, err error)
+}
+
 // Config的配置，有两个种用法
 // 一种是声明一个全局的配置，后面不停使用。
 // 另外一种是局部声明一个配置，然后使用WithXXX函数设置配置

conn.go

@@ -208,10 +208,11 @@ func (c *Conn) readDataFromNet(headArray *[enum.MaxFrameHeaderSize]byte, bufioPa
 		}
 	} else {
 		r := io.Reader(c.br)
+		var lr io.Reader
 		if c.readMaxMessage > 0 {
-			r = limitreader.NewLimitReader(c.br, c.readMaxMessage)
+			lr = limitreader.NewLimitReader(c.br, c.readMaxMessage)
 		}
-		f, err = frame.ReadFrameFromReaderV2(r, headArray, bufioPayload)
+		f, err = frame.ReadFrameFromReaderV3(r, lr, headArray, bufioPayload)
 	}
 	if err != nil {
 		c.writeAndMaybeOnClose(err)

conn_test.go

@@ -608,7 +608,7 @@ func TestFragmentFrame(t *testing.T) {
 		select {
 		case <-data:
 			atomic.AddInt32(&run, 1)
-		case <-time.After(500 * time.Hour):
+		case <-time.After(500 * time.Millisecond):
 		}
 
 		if atomic.LoadInt32(&run) != 1 {

go.mod

@@ -3,7 +3,7 @@ module github.com/antlabs/quickws
 go 1.21
 
 require (
-	github.com/antlabs/wsutil v0.1.10
+	github.com/antlabs/wsutil v0.1.11
 	golang.org/x/net v0.23.0
 )
 

go.sum

@@ -1,5 +1,5 @@
-github.com/antlabs/wsutil v0.1.10 h1:86p67dG8/iiQ+yZrHVl73OPHGnXfXopFSU0w84fLOdE=
-github.com/antlabs/wsutil v0.1.10/go.mod h1:Pk7xYOw3o5iEB6ukiOu+2uJMLYeMVVjJLazFD3okI2A=
+github.com/antlabs/wsutil v0.1.11 h1:bIVZ3Hxdq5ByZKu5OXL/cMtanEw6YlxdtUDiySI77Q0=
+github.com/antlabs/wsutil v0.1.11/go.mod h1:Pk7xYOw3o5iEB6ukiOu+2uJMLYeMVVjJLazFD3okI2A=
 github.com/klauspost/compress v1.17.8 h1:YcnTYrq7MikUT7k0Yb5eceMmALQPYBW/Xltxn0NAMnU=
 github.com/klauspost/compress v1.17.8/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
 golang.org/x/net v0.23.0 h1:7EYJ93RZ9vYSZAIb2x3lnuvqO5zneoD6IvWjuhfxjTs=

proxy.go

@@ -20,31 +20,33 @@ import (
 	"net"
 	"net/http"
 	"net/url"
+	"time"
 
 	"github.com/antlabs/wsutil/hostname"
 )
 
 type (
-	dialFunc  func(network, addr string) (c net.Conn, err error)
+	dialFunc  func(network, addr string, timeout time.Duration) (c net.Conn, err error)
 	httpProxy struct {
-		proxyAddr *url.URL
-		dial      func(network, addr string) (c net.Conn, err error)
+		proxyAddr   *url.URL
+		dialTimeout func(network, addr string, timeout time.Duration) (c net.Conn, err error)
+		timeout     time.Duration
 	}
 )
 
-var _ Dialer = (*httpProxy)(nil)
+var _ DialerTimeout = (*httpProxy)(nil)
 
 func newhttpProxy(u *url.URL, dial dialFunc) *httpProxy {
-	return &httpProxy{proxyAddr: u, dial: dial}
+	return &httpProxy{proxyAddr: u, dialTimeout: dial}
 }
 
-func (h *httpProxy) Dial(network, addr string) (c net.Conn, err error) {
+func (h *httpProxy) DialTimeout(network, addr string, timeout time.Duration) (c net.Conn, err error) {
 	if h.proxyAddr == nil {
-		return h.dial(network, addr)
+		return h.dialTimeout(network, addr, h.timeout)
 	}
 
 	hostName := hostname.GetHostName(h.proxyAddr)
-	c, err = h.dial(network, hostName)
+	c, err = h.dialTimeout(network, hostName, h.timeout)
 	if err != nil {
 		return nil, err
 	}

proxy_test.go

@@ -22,6 +22,7 @@ import (
 	"net/url"
 	"strings"
 	"testing"
+	"time"
 )
 
 type testServer struct {
@@ -133,7 +134,7 @@ func Test_Proxy(t *testing.T) {
 func Test_httpProxy_Dial(t *testing.T) {
 	type fields struct {
 		proxyAddr *url.URL
-		dial      func(network, addr string) (c net.Conn, err error)
+		dial      func(network, addr string, timeout time.Duration) (c net.Conn, err error)
 	}
 	type args struct {
 		network string
@@ -146,12 +147,12 @@ func Test_httpProxy_Dial(t *testing.T) {
 		wantC   net.Conn
 		wantErr bool
 	}{
-		// TODO: Add test cases.
+		// 0
 		{
 			name: "No proxy address",
 			fields: fields{
 				proxyAddr: nil,
-				dial: func(network, addr string) (c net.Conn, err error) {
+				dial: func(network, addr string, timeout time.Duration) (c net.Conn, err error) {
 					// Simulate successful dialing
 					return &net.TCPConn{}, errors.New("fail")
 				},
@@ -163,11 +164,12 @@ func Test_httpProxy_Dial(t *testing.T) {
 			wantC:   &net.TCPConn{},
 			wantErr: true,
 		},
+		// 1
 		{
 			name: "Proxy address",
 			fields: fields{
 				proxyAddr: &url.URL{Host: "1.2.3:8080", User: url.UserPassword("user", "password")},
-				dial: func(network, addr string) (c net.Conn, err error) {
+				dial: func(network, addr string, timeout time.Duration) (c net.Conn, err error) {
 					// Simulate successful dialing
 					return &net.TCPConn{}, errors.New("fail")
 				},
@@ -179,11 +181,12 @@ func Test_httpProxy_Dial(t *testing.T) {
 			wantC:   &net.TCPConn{},
 			wantErr: true,
 		},
+		// 2
 		{
 			name: "Proxy address",
 			fields: fields{
 				proxyAddr: &url.URL{Host: "1.2.3:8080", User: url.UserPassword("user", "password")},
-				dial: func(network, addr string) (c net.Conn, err error) {
+				dial: func(network, addr string, timeout time.Duration) (c net.Conn, err error) {
 					// Simulate successful dialing
 					return &net.TCPConn{}, nil
 				},
@@ -193,17 +196,16 @@ func Test_httpProxy_Dial(t *testing.T) {
 				addr:    "a.b.c:80",
 			},
 			wantC:   &net.TCPConn{},
-			wantErr: true,
+			wantErr: false,
 		},
 	}
 	for i, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			h := &httpProxy{
-				proxyAddr: tt.fields.proxyAddr,
-				dial:      tt.fields.dial,
+				proxyAddr:   tt.fields.proxyAddr,
+				dialTimeout: tt.fields.dial,
 			}
-			_, err := h.Dial(tt.args.network, tt.args.addr)
-			// gotC, err := h.Dial(tt.args.network, tt.args.addr)
+			_, err := h.dialTimeout(tt.args.network, tt.args.addr, 0)
 			if (err != nil) != tt.wantErr {
 				t.Errorf("index:%d, httpProxy.Dial() error = %v, wantErr %v", i, err, tt.wantErr)
 				return