Add UDP Like transport mode to snowflake

server/lib/http.go

@@ -108,6 +108,16 @@ func (handler *httpHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	// Pass the address of client as the remote address of incoming connection
 	clientIPParam := r.URL.Query().Get("client_ip")
 	addr := clientAddr(clientIPParam)
+	clientTransport := r.URL.Query().Get("protocol")
+
+	if clientTransport == "u" {
+		err = handler.turboTunnelUDPLikeMode(conn, addr)
+		if err != nil && err != io.EOF {
+			log.Println(err)
+			return
+		}
+		return
+	}
 
 	var token [len(turbotunnel.Token)]byte
 	_, err = io.ReadFull(conn, token[:])
@@ -221,6 +231,61 @@ func (handler *httpHandler) turbotunnelMode(conn net.Conn, addr net.Addr) error
 	return nil
 }
 
+func (handler *httpHandler) turboTunnelUDPLikeMode(conn net.Conn, addr net.Addr) error {
+	packetConnIDCon := packetConnIDConnServer{Conn: conn}
+	var packet [1600]byte
+	n, err := packetConnIDCon.Read(packet[:])
+	if err != nil {
+		return fmt.Errorf("reading ClientID: %v", err)
+	}
+	clientID, err := packetConnIDCon.GetClientID()
+	if err != nil {
+		return fmt.Errorf("reading ClientID: %v", err)
+	}
+	clientIDAddrMap.Set(clientID, addr)
+
+	pconn := handler.lookupPacketConn(clientID)
+	pconn.QueueIncoming(packet[:n], clientID)
+	var wg sync.WaitGroup
+	wg.Add(2)
+	done := make(chan struct{})
+	go func() {
+		defer wg.Done()
+		defer close(done) // Signal the write loop to finish
+		for {
+			n, err := packetConnIDCon.Read(packet[:])
+			if err != nil {
+				log.Println(err)
+				return
+			}
+			pconn.QueueIncoming(packet[:n], clientID)
+		}
+	}()
+	go func() {
+		defer wg.Done()
+		defer conn.Close() // Signal the read loop to finish
+		for {
+			select {
+			case <-done:
+				return
+			case p, ok := <-pconn.OutgoingQueue(clientID):
+				if !ok {
+					return
+				}
+				_, err := packetConnIDCon.Write(p)
+				pconn.Restore(p)
+				if err != nil {
+					log.Println(err)
+					return
+				}
+			}
+		}
+	}()
+
+	wg.Wait()
+	return nil
+}
+
 // ClientMapAddr is a string that represents a connecting client.
 type ClientMapAddr string
 

server/lib/packetIDConnServer.go

@@ -0,0 +1,52 @@
+package snowflake_server
+
+import (
+	"errors"
+	"net"
+
+	"gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/snowflake/v2/common/turbotunnel"
+)
+
+type ConnID = turbotunnel.ClientID
+
+type packetConnIDConnServer struct {
+	// This net.Conn must preserve message boundaries.
+	net.Conn
+	connID           ConnID
+	clientIDReceived bool
+}
+
+var ErrClientIDNotReceived = errors.New("ClientID not received")
+
+func (p *packetConnIDConnServer) GetClientID() (ConnID, error) {
+	if !p.clientIDReceived {
+		return p.connID, ErrClientIDNotReceived
+	}
+	return p.connID, nil
+}
+
+func (p *packetConnIDConnServer) Read(buf []byte) (n int, err error) {
+	n, err = p.Conn.Read(buf)
+	if err != nil {
+		return
+	}
+	switch buf[0] {
+	case 0xfe:
+		p.clientIDReceived = true
+		copy(p.connID[:], buf[1:9])
+		copy(buf[0:], buf[9:])
+		return n - 9, nil
+	case 0xff:
+		copy(buf[0:], buf[1:])
+		return n - 1, nil
+	}
+	return 0, nil
+}
+
+func (p *packetConnIDConnServer) Write(buf []byte) (n int, err error) {
+	n, err = p.Conn.Write(append([]byte{0xff}, buf...))
+	if err != nil {
+		return 0, err
+	}
+	return len(buf) - 1, nil
+}

server/lib/snowflake.go

@@ -253,7 +253,7 @@ func (l *SnowflakeListener) acceptSessions(ln *kcp.Listener) error {
 			return err
 		}
 		// Permit coalescing the payloads of consecutive sends.
-		conn.SetStreamMode(true)
+		conn.SetStreamMode(false)
 		// Set the maximum send and receive window sizes to a high number
 		// Removes KCP bottlenecks: https://gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/snowflake/-/issues/40026
 		conn.SetWindowSize(WindowSize, WindowSize)