From 1ba58c864d5b288047b7dbaa17b6369c4e77d1af Mon Sep 17 00:00:00 2001
From: Shelikhoo <xiaokangwang@outlook.com>
Date: Thu, 30 Jan 2025 13:35:36 +0000
Subject: [PATCH] refactor message boundary constrains

---
 client/lib/connwrapper.go             | 10 ----------
 client/lib/snowflake.go               |  3 +--
 client/lib/webrtc.go                  |  2 ++
 common/packetpadding/conn.go          | 15 +++------------
 common/websocketconn/websocketconn.go |  2 ++
 server/lib/http.go                    |  5 ++---
 6 files changed, 10 insertions(+), 27 deletions(-)

diff --git a/client/lib/connwrapper.go b/client/lib/connwrapper.go
index 0db9324..f8d0614 100644
--- a/client/lib/connwrapper.go
+++ b/client/lib/connwrapper.go
@@ -12,16 +12,6 @@ type ReadWriteCloserPreservesBoundary interface {
 	MessageBoundaryPreserved()
 }
 
-func ConfirmsReadWriteCloserPreservesMessageBoundary(rwc io.ReadWriteCloser) ReadWriteCloserPreservesBoundary {
-	return &messageBoundaryPreservedReadWriteCloser{rwc}
-}
-
-type messageBoundaryPreservedReadWriteCloser struct {
-	io.ReadWriteCloser
-}
-
-func (m *messageBoundaryPreservedReadWriteCloser) MessageBoundaryPreserved() {}
-
 var errENOSYS = errors.New("not implemented")
 
 func newPacketConnWrapper(localAddr, remoteAddr net.Addr, rwc ReadWriteCloserPreservesBoundary) net.PacketConn {
diff --git a/client/lib/snowflake.go b/client/lib/snowflake.go
index 92c36b4..0ede459 100644
--- a/client/lib/snowflake.go
+++ b/client/lib/snowflake.go
@@ -344,8 +344,7 @@ func newSession(snowflakes SnowflakeCollector) (net.PacketConn, *smux.Session, e
 		log.Println("---- Handler: snowflake assigned ----")
 
 		packetConnWrapper := newPacketConnWrapper(dummyAddr{}, dummyAddr{},
-			packetpadding.NewPaddableConnection(
-				ConfirmsReadWriteCloserPreservesMessageBoundary(conn),
+			packetpadding.NewPaddableConnection(conn,
 				packetpadding.New()))
 
 		return packetConnWrapper, nil
diff --git a/client/lib/webrtc.go b/client/lib/webrtc.go
index 7af1186..ed745d9 100644
--- a/client/lib/webrtc.go
+++ b/client/lib/webrtc.go
@@ -412,3 +412,5 @@ func (c *WebRTCPeer) cleanup() {
 		}
 	}
 }
+
+func (c *WebRTCPeer) MessageBoundaryPreserved() {}
diff --git a/common/packetpadding/conn.go b/common/packetpadding/conn.go
index 8993745..b214e70 100644
--- a/common/packetpadding/conn.go
+++ b/common/packetpadding/conn.go
@@ -1,23 +1,14 @@
 package packetpadding
 
-import "io"
+import (
+	"io"
+)
 
 type ReadWriteCloserPreservesBoundary interface {
 	io.ReadWriteCloser
 	MessageBoundaryPreserved()
 }
 
-type messageBoundaryPreservedReadWriteCloser struct {
-	io.ReadWriteCloser
-}
-
-func (m *messageBoundaryPreservedReadWriteCloser) MessageBoundaryPreserved() {
-}
-
-func ConfirmsReadWriteCloserPreservesMessageBoundary(rwc io.ReadWriteCloser) ReadWriteCloserPreservesBoundary {
-	return &messageBoundaryPreservedReadWriteCloser{rwc}
-}
-
 type PaddableConnection interface {
 	ReadWriteCloserPreservesBoundary
 }
diff --git a/common/websocketconn/websocketconn.go b/common/websocketconn/websocketconn.go
index e5256df..0178f3d 100644
--- a/common/websocketconn/websocketconn.go
+++ b/common/websocketconn/websocketconn.go
@@ -41,6 +41,8 @@ func (conn *Conn) SetDeadline(t time.Time) error {
 	return err
 }
 
+func (conn *Conn) MessageBoundaryPreserved() {}
+
 func readLoop(w io.Writer, ws *websocket.Conn) error {
 	var buf [2048]byte
 	for {
diff --git a/server/lib/http.go b/server/lib/http.go
index 39eac9b..89c8a6b 100644
--- a/server/lib/http.go
+++ b/server/lib/http.go
@@ -115,7 +115,7 @@ func (handler *httpHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
-func (handler *httpHandler) turboTunnelUDPLikeMode(conn net.Conn, addr net.Addr, protocol string) error {
+func (handler *httpHandler) turboTunnelUDPLikeMode(conn *websocketconn.Conn, addr net.Addr, protocol string) error {
 	// Read the ClientID from the WebRTC data channel protocol string. Every
 	// packet received on this WebSocket connection pertains to the same
 	// ClientID.
@@ -143,8 +143,7 @@ func (handler *httpHandler) turboTunnelUDPLikeMode(conn net.Conn, addr net.Addr,
 	wg.Add(2)
 	done := make(chan struct{})
 
-	connPaddable := packetpadding.NewPaddableConnection(
-		packetpadding.ConfirmsReadWriteCloserPreservesMessageBoundary(conn), packetpadding.New())
+	connPaddable := packetpadding.NewPaddableConnection(conn, packetpadding.New())
 
 	// The remainder of the WebSocket stream consists of packets, one packet
 	// per WebSocket message. We read them one by one and feed them into the