隐藏conn.Write接口

Author: guonaihong

autobahn/server/autobahn-server.go

@@ -27,6 +27,10 @@ var keyPEMBlock []byte
 type echoHandler struct{}
 
 func (e *echoHandler) OnOpen(c *greatws.Conn) {
+	// err := c.WriteMessage(greatws.Binary, make([]byte, 1<<28))
+	// if err != nil {
+	// 	fmt.Printf("%s\n", err)
+	// }
 	// fmt.Printf("OnOpen: %p\n", c)
 }
 
@@ -132,7 +136,7 @@ func (h *handler) echoRunStream2(w http.ResponseWriter, r *http.Request) {
 		greatws.WithServerIgnorePong(),
 		greatws.WithServerCallback(&echoHandler{}),
 		greatws.WithServerEnableUTF8Check(),
-		greatws.WithServerReadTimeout(5 * time.Second),
+		// greatws.WithServerReadTimeout(5 * time.Second),
 		greatws.WithServerMultiEventLoop(h.m),
 		greatws.WithServerStreamMode(),
 		greatws.WithServerCallbackInEventLoop(),

conn_core.go

@@ -581,7 +581,7 @@ func (c *Conn) WriteMessage(op Opcode, writeBuf []byte) (err error) {
 	var fw fixedwriter.FixedWriter
 
 	c.mu.Lock()
-	err = frame.WriteFrame(&fw, c, writeBuf, true, rsv1, c.client, op, maskValue)
+	err = frame.WriteFrame(&fw, connToNewConn(c), writeBuf, true, rsv1, c.client, op, maskValue)
 	c.mu.Unlock()
 
 	return err
@@ -619,14 +619,14 @@ func (c *Conn) writeFragment(op Opcode, writeBuf []byte, maxFragment int /*单
 	var fw fixedwriter.FixedWriter
 	for len(writeBuf) > 0 {
 		if len(writeBuf) > maxFragment {
-			if err := frame.WriteFrame(&fw, c, writeBuf[:maxFragment], false, rsv1, c.client, op, maskValue); err != nil {
+			if err := frame.WriteFrame(&fw, connToNewConn(c), writeBuf[:maxFragment], false, rsv1, c.client, op, maskValue); err != nil {
 				return err
 			}
 			writeBuf = writeBuf[maxFragment:]
 			op = Continuation
 			continue
 		}
-		return frame.WriteFrame(&fw, c, writeBuf, true, rsv1, c.client, op, maskValue)
+		return frame.WriteFrame(&fw, connToNewConn(c), writeBuf, true, rsv1, c.client, op, maskValue)
 	}
 	return nil
 }

conn_test.go

@@ -253,7 +253,7 @@ func Test_ReadMessage(t *testing.T) {
 		// err = con.WriteMessage(Binary, []byte("hello"))
 		maskValue := rand.Uint32()
 		var fw fixedwriter.FixedWriter
-		err = frame.WriteFrame(&fw, con, []byte("hello"), true, true, con.client, Binary, maskValue)
+		err = frame.WriteFrame(&fw, connToNewConn(con), []byte("hello"), true, true, con.client, Binary, maskValue)
 		if err != nil {
 			t.Error(err)
 		}
@@ -314,7 +314,7 @@ func Test_ReadMessage(t *testing.T) {
 		// err = con.WriteMessage(Binary, []byte("hello"))
 		maskValue := rand.Uint32()
 		var fw fixedwriter.FixedWriter
-		err = frame.WriteFrame(&fw, con, []byte("hello"), true, true, con.client, Ping, maskValue)
+		err = frame.WriteFrame(&fw, connToNewConn(con), []byte("hello"), true, true, con.client, Ping, maskValue)
 		if err != nil {
 			t.Error(err)
 		}
@@ -461,12 +461,12 @@ func TestFragmentFrame(t *testing.T) {
 
 		maskValue := rand.Uint32()
 		var fw fixedwriter.FixedWriter
-		err = frame.WriteFrame(&fw, con, []byte("h"), false, false, con.client, Text, maskValue)
+		err = frame.WriteFrame(&fw, connToNewConn(con), []byte("h"), false, false, con.client, Text, maskValue)
 		if err != nil {
 			t.Error(err)
 		}
 		maskValue = rand.Uint32()
-		err = frame.WriteFrame(&fw, con, []byte{}, true, false, con.client, Text, maskValue)
+		err = frame.WriteFrame(&fw, connToNewConn(con), []byte{}, true, false, con.client, Text, maskValue)
 		if err != nil {
 			t.Error(err)
 		}

conn_unix.go

@@ -185,7 +185,9 @@ func (c *Conn) getPtr() int {
 	return int(uintptr(unsafe.Pointer(c)))
 }
 
-func (c *Conn) Write(b []byte) (n int, err error) {
+// Conn Write入口， 原始定义是func (c *Conn) Write() (n int, err error)
+// 会引起误用，所以隐藏起来, 作为一个websocket库，直接暴露tcp的write接口也不合适
+func connWrite(c *Conn, b []byte) (n int, err error) {
 	if c.isClosed() {
 		return 0, ErrClosed
 	}

conn_write.go

@@ -0,0 +1,31 @@
+// Copyright 2023-2024 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.
+
+//go:build linux || darwin || netbsd || freebsd || openbsd || dragonfly
+// +build linux darwin netbsd freebsd openbsd dragonfly
+
+package greatws
+
+import "unsafe"
+
+type newConnWrite Conn
+
+func connToNewConn(c *Conn) *newConnWrite {
+	return (*newConnWrite)(unsafe.Pointer(c))
+}
+
+func (c *newConnWrite) Write(p []byte) (n int, err error) {
+	c2 := (*Conn)(unsafe.Pointer(c))
+	return connWrite(c2, p)
+}

multi_event_loops.go

@@ -186,6 +186,9 @@ func (m *MultiEventLoop) getEventLoop(fd int) *EventLoop {
 // 添加一个连接到多路事件循环
 func (m *MultiEventLoop) add(c *Conn) error {
 	fd := c.getFd()
+	if fd == -1 {
+		return nil
+	}
 	index := fd % len(m.loops)
 	m.safeConns.addConn(c)
 	// m.loops[index].conns.Store(fd, c)