提升测试覆盖度

benchmark_rand_test.go

@@ -1,3 +1,16 @@
+// Copyright 2021-2023 antlabs. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
 package quickws
 
 import (

client.go

@@ -43,7 +43,7 @@ type DialOption struct {
 	Config
 }
 
-func ClientOpt(opts ...ClientOption) *Config {
+func ClientOptionToConf(opts ...ClientOption) *Config {
 	var dial DialOption
 	dial.defaultSetting()
 	for _, o := range opts {
@@ -52,7 +52,7 @@ func ClientOpt(opts ...ClientOption) *Config {
 	return &dial.Config
 }
 
-func DialWithConfig(rawUrl string, conf *Config) (*Conn, error) {
+func DialConf(rawUrl string, conf *Config) (*Conn, error) {
 	var dial DialOption
 	u, err := url.Parse(rawUrl)
 	if err != nil {
@@ -99,7 +99,7 @@ func (d *DialOption) handshake() (*http.Request, string, error) {
 		d.u.Scheme = "http"
 	default:
 		// TODO 返回错误
-		return nil, "", fmt.Errorf("未知的scheme:%s", d.u.Scheme)
+		return nil, "", fmt.Errorf("Unknown scheme, only supports ws:// or wss://: got %s", d.u.Scheme)
 	}
 
 	// 满足4.1

client_test.go

@@ -0,0 +1,143 @@
+// Copyright 2021-2023 antlabs. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package quickws
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync/atomic"
+	"testing"
+)
+
+// 测试客户端Dial
+func Test_Client_Dial(t *testing.T) {
+	t.Run("Dial: valid resp: status code fail", func(t *testing.T) {
+		done := make(chan bool, 1)
+		run := int32(0)
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			atomic.AddInt32(&run, int32(1))
+			done <- true
+		}))
+
+		defer ts.Close()
+
+		rawURL := strings.ReplaceAll(ts.URL, "http", "ws")
+		_, err := Dial(rawURL)
+		if err == nil {
+			t.Fatal("should be error")
+		}
+	})
+
+	t.Run("DialConf: valid resp : status code fail", func(t *testing.T) {
+		done := make(chan bool, 1)
+		run := int32(0)
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			atomic.AddInt32(&run, int32(1))
+			done <- true
+		}))
+
+		defer ts.Close()
+
+		cnf := ClientOptionToConf()
+		rawURL := strings.ReplaceAll(ts.URL, "http", "ws")
+		_, err := DialConf(rawURL, cnf)
+		if err == nil {
+			t.Fatal("should be error")
+		}
+	})
+
+	t.Run("Dial: valid resp: Upgrade field fail", func(t *testing.T) {
+		done := make(chan bool, 1)
+		run := int32(0)
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			atomic.AddInt32(&run, int32(1))
+			w.WriteHeader(101)
+			w.Header().Set("Upgrade", "xx")
+			done <- true
+		}))
+
+		defer ts.Close()
+
+		rawURL := strings.ReplaceAll(ts.URL, "http", "ws")
+		_, err := Dial(rawURL)
+		if err == nil {
+			t.Fatal("should be error")
+		}
+	})
+
+	t.Run("DialConf: valid resp: Upgrade field fail", func(t *testing.T) {
+		done := make(chan bool, 1)
+		run := int32(0)
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			atomic.AddInt32(&run, int32(1))
+			w.WriteHeader(101)
+			w.Header().Set("Upgrade", "xx")
+			done <- true
+		}))
+
+		defer ts.Close()
+
+		cnf := ClientOptionToConf()
+		rawURL := strings.ReplaceAll(ts.URL, "http", "ws")
+		_, err := DialConf(rawURL, cnf)
+		if err == nil {
+			t.Fatal("should be error")
+		}
+	})
+
+	t.Run("Dial: valid resp: Connection fail", func(t *testing.T) {
+		done := make(chan bool, 1)
+		run := int32(0)
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			atomic.AddInt32(&run, int32(1))
+			w.Header().Set("Upgrade", "websocket")
+			w.Header().Set("Connection", "xx")
+			w.WriteHeader(101)
+			done <- true
+		}))
+
+		defer ts.Close()
+
+		rawURL := strings.ReplaceAll(ts.URL, "http", "ws")
+		_, err := Dial(rawURL)
+		if err == nil {
+			t.Fatal("should be error")
+		}
+	})
+
+	t.Run("DialConf: valid resp: Connection fail", func(t *testing.T) {
+		done := make(chan bool, 1)
+		run := int32(0)
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			atomic.AddInt32(&run, int32(1))
+			w.Header().Set("Upgrade", "websocket")
+			w.Header().Set("Connection", "xx")
+			w.WriteHeader(101)
+			done <- true
+		}))
+
+		defer ts.Close()
+
+		cnf := ClientOptionToConf()
+		rawURL := strings.ReplaceAll(ts.URL, "http", "ws")
+		_, err := DialConf(rawURL, cnf)
+		if err == nil {
+			t.Fatal("should be error")
+		} else {
+			fmt.Printf("err: %v\n", err)
+		}
+	})
+}

message_test.go

@@ -1,3 +1,16 @@
+// Copyright 2021-2023 antlabs. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
 package quickws
 
 import (

parse_mode.go

@@ -1,3 +1,16 @@
+// Copyright 2021-2023 antlabs. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
 package quickws
 
 type parseMode int32

server_handshake.go

@@ -23,28 +23,17 @@ import (
 )
 
 var (
-	ErrNotFoundHijacker     = errors.New("not found Hijacker")
-	bytesHeaderUpgrade      = []byte("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n")
-	bytesSecWebSocketAccept = []byte("Sec-WebSocket-Accept")
-	bytesHeaderExtensions   = []byte("Sec-WebSocket-Extensions: permessage-deflate; server_no_context_takeover; client_no_context_takeover\r\n")
-	bytesCRLF               = []byte("\r\n")
-	bytesColon              = []byte(": ")
+	ErrNotFoundHijacker   = errors.New("not found Hijacker")
+	bytesHeaderUpgrade    = []byte("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept:")
+	bytesHeaderExtensions = []byte("Sec-WebSocket-Extensions: permessage-deflate; server_no_context_takeover; client_no_context_takeover\r\n")
+	bytesCRLF             = []byte("\r\n")
+	bytesColon            = []byte(": ")
 )
 
 type ConnOption struct {
 	Config
 }
 
-func writeHeaderKey(w io.Writer, key []byte) (err error) {
-	if _, err = w.Write(key); err != nil {
-		return
-	}
-	if _, err = w.Write(bytesColon); err != nil {
-		return
-	}
-	return
-}
-
 func writeHeaderVal(w io.Writer, val []byte) (err error) {
 	if _, err = w.Write(val); err != nil {
 		return
@@ -59,14 +48,12 @@ func writeHeaderVal(w io.Writer, val []byte) (err error) {
 // https://datatracker.ietf.org/doc/html/rfc6455#section-4.2.2
 // 第5小点
 func prepareWriteResponse(r *http.Request, w io.Writer, cnf *Config) (err error) {
+	// 写入响应头
+	// 写入Sec-WebSocket-Accept key
 	if _, err = w.Write(bytesHeaderUpgrade); err != nil {
 		return
 	}
 
-	// 写入Sec-WebSocket-Accept key
-	if err = writeHeaderKey(w, bytesSecWebSocketAccept); err != nil {
-		return
-	}
 	v := secWebSocketAcceptVal(r.Header.Get("Sec-WebSocket-Key"))
 	// 写入Sec-WebSocket-Accept vla
 	if err = writeHeaderVal(w, StringToBytes(v)); err != nil {

quickws_server_test.go → server_test.go

@@ -1,3 +1,16 @@
+// Copyright 2021-2023 antlabs. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
 package quickws
 
 import (

upgrade.go

@@ -19,6 +19,7 @@ import (
 	"bytes"
 	"net"
 	"net/http"
+	"time"
 
 	"github.com/antlabs/wsutil/bufio2"
 	"github.com/antlabs/wsutil/bytespool"
@@ -112,5 +113,7 @@ func upgradeInner(w http.ResponseWriter, r *http.Request, conf *Config) (c *Conn
 	if conf.parseMode == ParseModeWindows {
 		fr.Init(conn, bytespool.GetBytes(conf.initPayloadSize()+enum.MaxFrameHeaderSize))
 	}
+
+	conn.SetDeadline(time.Time{})
 	return newConn(conn, false, conf, fr, read, bp), nil
 }