sing-box/transport/v2raywebsocket/conn.go

235 lines
5 KiB
Go
Raw Normal View History

2022-08-22 12:20:56 +00:00
package v2raywebsocket
import (
2022-08-26 02:22:20 +00:00
"context"
2022-08-22 12:20:56 +00:00
"encoding/base64"
"io"
"net"
"net/http"
"os"
"time"
2022-08-25 02:44:58 +00:00
C "github.com/sagernet/sing-box/constant"
2022-10-07 12:18:50 +00:00
"github.com/sagernet/sing/common"
2022-09-23 09:21:24 +00:00
"github.com/sagernet/sing/common/buf"
2022-08-22 12:20:56 +00:00
E "github.com/sagernet/sing/common/exceptions"
2022-09-13 02:41:10 +00:00
"github.com/sagernet/websocket"
2022-08-22 12:20:56 +00:00
)
type WebsocketConn struct {
*websocket.Conn
2022-09-23 09:21:24 +00:00
*Writer
2022-08-22 12:20:56 +00:00
remoteAddr net.Addr
reader io.Reader
}
2022-09-23 09:21:24 +00:00
func NewServerConn(wsConn *websocket.Conn, remoteAddr net.Addr) *WebsocketConn {
return &WebsocketConn{
Conn: wsConn,
remoteAddr: remoteAddr,
2022-09-30 03:27:18 +00:00
Writer: NewWriter(wsConn, true),
2022-09-23 09:21:24 +00:00
}
}
2022-08-25 02:44:58 +00:00
func (c *WebsocketConn) Close() error {
err := c.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(C.TCPTimeout))
if err != nil {
return c.Conn.Close()
}
return nil
}
2022-08-22 12:20:56 +00:00
func (c *WebsocketConn) Read(b []byte) (n int, err error) {
for {
if c.reader == nil {
_, c.reader, err = c.NextReader()
if err != nil {
2022-08-25 02:44:58 +00:00
err = wrapError(err)
2022-08-22 12:20:56 +00:00
return
}
}
n, err = c.reader.Read(b)
if E.IsMulti(err, io.EOF) {
c.reader = nil
continue
}
2022-08-25 02:44:58 +00:00
err = wrapError(err)
2022-08-22 12:20:56 +00:00
return
}
}
func (c *WebsocketConn) RemoteAddr() net.Addr {
if c.remoteAddr != nil {
return c.remoteAddr
}
return c.Conn.RemoteAddr()
}
func (c *WebsocketConn) SetDeadline(t time.Time) error {
return os.ErrInvalid
}
2022-10-07 12:18:50 +00:00
func (c *WebsocketConn) Upstream() any {
return c.Conn.NetConn()
}
func (c *WebsocketConn) UpstreamWriter() any {
return c.Writer
2022-09-23 09:21:24 +00:00
}
2022-08-22 12:20:56 +00:00
type EarlyWebsocketConn struct {
*Client
2022-08-26 02:22:20 +00:00
ctx context.Context
2022-08-22 12:20:56 +00:00
conn *WebsocketConn
create chan struct{}
}
func (c *EarlyWebsocketConn) Read(b []byte) (n int, err error) {
if c.conn == nil {
<-c.create
}
return c.conn.Read(b)
}
func (c *EarlyWebsocketConn) Write(b []byte) (n int, err error) {
if c.conn != nil {
return c.conn.Write(b)
}
var (
earlyData []byte
lateData []byte
conn *websocket.Conn
response *http.Response
)
2022-09-23 09:21:24 +00:00
if len(b) > int(c.maxEarlyData) {
earlyData = b[:c.maxEarlyData]
lateData = b[c.maxEarlyData:]
2022-08-22 12:20:56 +00:00
} else {
earlyData = b
}
if len(earlyData) > 0 {
earlyDataString := base64.RawURLEncoding.EncodeToString(earlyData)
if c.earlyDataHeaderName == "" {
2022-08-26 02:22:20 +00:00
conn, response, err = c.dialer.DialContext(c.ctx, c.uri+earlyDataString, c.headers)
2022-08-22 12:20:56 +00:00
} else {
headers := c.headers.Clone()
headers.Set(c.earlyDataHeaderName, earlyDataString)
2022-08-26 02:22:20 +00:00
conn, response, err = c.dialer.DialContext(c.ctx, c.uri, headers)
2022-08-22 12:20:56 +00:00
}
} else {
2022-08-26 02:22:20 +00:00
conn, response, err = c.dialer.DialContext(c.ctx, c.uri, c.headers)
2022-08-22 12:20:56 +00:00
}
if err != nil {
return 0, wrapDialError(response, err)
}
2022-09-30 03:27:18 +00:00
c.conn = &WebsocketConn{Conn: conn, Writer: NewWriter(conn, false)}
2022-08-22 12:20:56 +00:00
close(c.create)
if len(lateData) > 0 {
_, err = c.conn.Write(lateData)
}
if err != nil {
return
}
return len(b), nil
}
2022-09-23 09:21:24 +00:00
func (c *EarlyWebsocketConn) WriteBuffer(buffer *buf.Buffer) error {
if c.conn != nil {
return c.conn.WriteBuffer(buffer)
}
var (
earlyData []byte
lateData []byte
conn *websocket.Conn
response *http.Response
err error
)
if buffer.Len() > int(c.maxEarlyData) {
earlyData = buffer.Bytes()[:c.maxEarlyData]
lateData = buffer.Bytes()[c.maxEarlyData:]
} else {
earlyData = buffer.Bytes()
}
if len(earlyData) > 0 {
earlyDataString := base64.RawURLEncoding.EncodeToString(earlyData)
if c.earlyDataHeaderName == "" {
conn, response, err = c.dialer.DialContext(c.ctx, c.uri+earlyDataString, c.headers)
} else {
headers := c.headers.Clone()
headers.Set(c.earlyDataHeaderName, earlyDataString)
conn, response, err = c.dialer.DialContext(c.ctx, c.uri, headers)
}
} else {
conn, response, err = c.dialer.DialContext(c.ctx, c.uri, c.headers)
}
if err != nil {
return wrapDialError(response, err)
}
2022-09-30 03:27:18 +00:00
c.conn = &WebsocketConn{Conn: conn, Writer: NewWriter(conn, false)}
2022-09-23 09:21:24 +00:00
close(c.create)
if len(lateData) > 0 {
_, err = c.conn.Write(lateData)
}
return err
}
2022-08-22 12:20:56 +00:00
func (c *EarlyWebsocketConn) Close() error {
if c.conn == nil {
return nil
}
return c.conn.Close()
}
func (c *EarlyWebsocketConn) LocalAddr() net.Addr {
if c.conn == nil {
return nil
}
return c.conn.LocalAddr()
}
func (c *EarlyWebsocketConn) RemoteAddr() net.Addr {
if c.conn == nil {
return nil
}
return c.conn.RemoteAddr()
}
func (c *EarlyWebsocketConn) SetDeadline(t time.Time) error {
if c.conn == nil {
return os.ErrInvalid
}
return c.conn.SetDeadline(t)
}
func (c *EarlyWebsocketConn) SetReadDeadline(t time.Time) error {
if c.conn == nil {
return os.ErrInvalid
}
return c.conn.SetReadDeadline(t)
}
func (c *EarlyWebsocketConn) SetWriteDeadline(t time.Time) error {
if c.conn == nil {
return os.ErrInvalid
}
return c.conn.SetWriteDeadline(t)
}
2022-08-25 02:44:58 +00:00
2022-10-07 12:18:50 +00:00
func (c *EarlyWebsocketConn) Upstream() any {
return common.PtrOrNil(c.conn)
}
func (c *EarlyWebsocketConn) LazyHeadroom() bool {
return c.conn == nil
2022-09-23 09:21:24 +00:00
}
2022-08-25 02:44:58 +00:00
func wrapError(err error) error {
if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
2022-09-01 12:04:04 +00:00
return io.EOF
}
if websocket.IsCloseError(err, websocket.CloseAbnormalClosure) {
2022-08-25 02:44:58 +00:00
return net.ErrClosed
}
return err
}