optimize BufferReset

lxzan · Sep 20, 2023 · 53f5dbd · 53f5dbd
1 parent 376c264
commit 53f5dbd
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 19 deletions.
diff --git a/internal/utils.go b/internal/utils.go
@@ -199,14 +199,6 @@ func ToBinaryNumber[T Integer](n T) T {
 	return x
 }
 
-type Buffer struct {
-	buf      []byte
-	off      int
-	lastRead int8
-}
-
-//go:nosplit
-func BufferReset(b *bytes.Buffer, p []byte) {
-	buffer := (*Buffer)(unsafe.Pointer(b)) // 类型强转
-	buffer.buf = p                         // 修改后面的属性一定要加偏移量!!!
-}
+// BufferReset 重置buffer底层切片
+// 修改后面的属性一定要加偏移量!!!
+func BufferReset(b *bytes.Buffer, p []byte) { *(*[]byte)(unsafe.Pointer(b)) = p }
diff --git a/internal/utils_test.go b/internal/utils_test.go
@@ -7,6 +7,7 @@ import (
 	"hash/fnv"
 	"io"
 	"net/http"
+	"reflect"
 	"strings"
 	"testing"
 
@@ -212,9 +213,35 @@ func TestGetIntersectionElem(t *testing.T) {
 }
 
 func TestResetBuffer(t *testing.T) {
-	var buf = bytes.NewBufferString("hello")
-	var p = buf.Bytes()
-	p = append(p, "world"...)
-	BufferReset(buf, p)
-	assert.Equal(t, "helloworld", buf.String())
+	{
+		var buf = bytes.NewBufferString("hello")
+		var p = AlphabetNumeric.Generate(1024)
+		BufferReset(buf, p)
+		assert.Equal(t, string(p), buf.String())
+	}
+
+	{
+		var buf = bytes.NewBufferString("")
+		BufferReset(buf, []byte("hello"))
+		assert.Equal(t, "hello", buf.String())
+	}
+
+	{
+		var buf = bytes.NewBufferString("hello")
+		buf.Next(2)
+		assert.Equal(t, "llo", buf.String())
+		BufferReset(buf, []byte("aha"))
+		assert.Equal(t, "a", buf.String())
+		assert.Equal(t, int64(2), reflect.ValueOf(buf).Elem().FieldByName("off").Int())
+		assert.Equal(t, int64(-1), reflect.ValueOf(buf).Elem().FieldByName("lastRead").Int())
+	}
+
+	{
+		var buf = bytes.NewBufferString("hello")
+		var p = buf.Bytes()
+		io.ReadAll(buf)
+		assert.Equal(t, "", buf.String())
+		BufferReset(buf, p)
+		assert.Equal(t, "hello", buf.String())
+	}
 }
diff --git a/option_test.go b/option_test.go
@@ -2,6 +2,7 @@ package gws
 
 import (
 	"compress/flate"
+	"net/http"
 	"testing"
 	"time"
 
@@ -49,7 +50,12 @@ func validateClientOption(as *assert.Assertions, option *ClientOption) {
 // 检查默认配置
 func TestDefaultUpgrader(t *testing.T) {
 	var as = assert.New(t)
-	var updrader = NewUpgrader(new(BuiltinEventHandler), nil)
+	var updrader = NewUpgrader(new(BuiltinEventHandler), &ServerOption{
+		ResponseHeader: http.Header{
+			"Sec-Websocket-Extensions": []string{"chat"},
+			"X-Server":                 []string{"gws"},
+		},
+	})
 	var config = updrader.option.getConfig()
 	as.Equal(false, config.CompressEnabled)
 	as.Equal(false, config.ReadAsyncEnabled)
@@ -66,6 +72,8 @@ func TestDefaultUpgrader(t *testing.T) {
 	as.NotNil(updrader.option.Authorize)
 	as.NotNil(updrader.option.NewSessionStorage)
 	as.Nil(updrader.option.SubProtocols)
+	as.Equal("", updrader.option.ResponseHeader.Get("Sec-Websocket-Extensions"))
+	as.Equal("gws", updrader.option.ResponseHeader.Get("X-Server"))
 	validateServerOption(as, updrader)
 }
 

diff --git a/reader.go b/reader.go
@@ -93,8 +93,7 @@ func (c *Conn) readMessage() error {
 
 	var fin = c.fh.GetFIN()
 	var buf, index = binaryPool.Get(contentLength)
-	var p = buf.Bytes()
-	p = p[:contentLength]
+	var p = buf.Bytes()[:contentLength]
 	if err := internal.ReadN(c.br, p); err != nil {
 		return err
 	}