admin
e6956aa001
- React frontend with route-level code splitting - Backend rebranded from Baihu to TaskPool - DB brand migration script and local compatibility
99 lines
2.3 KiB
Go
99 lines
2.3 KiB
Go
package tunnel
|
|
|
|
import (
|
|
"io"
|
|
"net"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
// Copyright (c) 2026 engigu (TaskPool). All rights reserved.
|
|
// Use of this source code is governed by the Apache License 2.0.
|
|
//
|
|
// 【重要声明 / IMPORTANT NOTICE】
|
|
// 本代码（包括其架构设计与核心实现）属于任务池（TaskPool）开源项目的一部分。
|
|
// 任何个人或组织在引用、移植、修改或重新分发此文件中的任何代码时，必须保留本版权声明，
|
|
// 并在您的衍生作品、文档、软件关于页面或说明文件中显式声明引用自任务池（TaskPool）。
|
|
//
|
|
// Anyone referencing, porting, modifying, or redistributing this code must retain this
|
|
// copyright notice and explicitly state the source: TaskPool (github.com/engigu/taskpool).
|
|
|
|
// wsConn 包装 gorilla websocket 以实现 net.Conn 接口
|
|
type wsConn struct {
|
|
conn *websocket.Conn
|
|
reader io.Reader
|
|
}
|
|
|
|
// NetConn 将 WebSocket 连接转换为标准 net.Conn
|
|
func NetConn(conn *websocket.Conn) net.Conn {
|
|
return &wsConn{
|
|
conn: conn,
|
|
}
|
|
}
|
|
|
|
func (c *wsConn) Read(b []byte) (n int, err error) {
|
|
for {
|
|
if c.reader == nil {
|
|
_, r, err := c.conn.NextReader()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
c.reader = r
|
|
}
|
|
n, err = c.reader.Read(b)
|
|
if err == io.EOF {
|
|
c.reader = nil
|
|
if n > 0 {
|
|
return n, nil
|
|
}
|
|
continue
|
|
}
|
|
return n, err
|
|
}
|
|
}
|
|
|
|
func (c *wsConn) Write(b []byte) (n int, err error) {
|
|
// Yamux 保证了对底层连接的序列化写入（不会出现并发写），
|
|
// 因此我们这里不需要加互斥锁 (Mutex)。
|
|
w, err := c.conn.NextWriter(websocket.BinaryMessage)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
n, err = w.Write(b)
|
|
|
|
// 必须关闭 writer 以触发 WebSocket 的消息帧封包发送
|
|
closeErr := w.Close()
|
|
if err == nil {
|
|
err = closeErr
|
|
}
|
|
return n, err
|
|
}
|
|
|
|
func (c *wsConn) Close() error {
|
|
return c.conn.Close()
|
|
}
|
|
|
|
func (c *wsConn) LocalAddr() net.Addr {
|
|
return c.conn.LocalAddr()
|
|
}
|
|
|
|
func (c *wsConn) RemoteAddr() net.Addr {
|
|
return c.conn.RemoteAddr()
|
|
}
|
|
|
|
func (c *wsConn) SetDeadline(t time.Time) error {
|
|
if err := c.conn.SetReadDeadline(t); err != nil {
|
|
return err
|
|
}
|
|
return c.conn.SetWriteDeadline(t)
|
|
}
|
|
|
|
func (c *wsConn) SetReadDeadline(t time.Time) error {
|
|
return c.conn.SetReadDeadline(t)
|
|
}
|
|
|
|
func (c *wsConn) SetWriteDeadline(t time.Time) error {
|
|
return c.conn.SetWriteDeadline(t)
|
|
}
|