duorameng
88 lines
1.6 KiB
Go
88 lines
1.6 KiB
Go
package tunnel
|
|
|
|
import (
|
|
"io"
|
|
"net"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
// 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)
|
|
}
|