initial commit
This commit is contained in:
ginuerzh
38
dialer/mws/conn.go
Normal file
38
dialer/mws/conn.go
Normal file
@ -0,0 +1,38 @@
|
||||
package mws
|
||||
|
||||
import (
|
||||
"net"
|
||||
|
||||
"github.com/xtaci/smux"
|
||||
)
|
||||
|
||||
type muxSession struct {
|
||||
conn net.Conn
|
||||
session *smux.Session
|
||||
}
|
||||
|
||||
func (session *muxSession) GetConn() (net.Conn, error) {
|
||||
return session.session.OpenStream()
|
||||
}
|
||||
|
||||
func (session *muxSession) Accept() (net.Conn, error) {
|
||||
return session.session.AcceptStream()
|
||||
}
|
||||
|
||||
func (session *muxSession) Close() error {
|
||||
if session.session == nil {
|
||||
return nil
|
||||
}
|
||||
return session.session.Close()
|
||||
}
|
||||
|
||||
func (session *muxSession) IsClosed() bool {
|
||||
if session.session == nil {
|
||||
return true
|
||||
}
|
||||
return session.session.IsClosed()
|
||||
}
|
||||
|
||||
func (session *muxSession) NumStreams() int {
|
||||
return session.session.NumStreams()
|
||||
}
|
||||
215
dialer/mws/dialer.go
Normal file
215
dialer/mws/dialer.go
Normal file
@ -0,0 +1,215 @@
|
||||
package mws
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net"
|
||||
"net/url"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/go-gost/gost/v3/pkg/dialer"
|
||||
md "github.com/go-gost/gost/v3/pkg/metadata"
|
||||
"github.com/go-gost/gost/v3/pkg/registry"
|
||||
ws_util "github.com/go-gost/x/internal/util/ws"
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/xtaci/smux"
|
||||
)
|
||||
|
||||
func init() {
|
||||
registry.DialerRegistry().Register("mws", NewDialer)
|
||||
registry.DialerRegistry().Register("mwss", NewTLSDialer)
|
||||
}
|
||||
|
||||
type mwsDialer struct {
|
||||
sessions map[string]*muxSession
|
||||
sessionMutex sync.Mutex
|
||||
tlsEnabled bool
|
||||
md metadata
|
||||
options dialer.Options
|
||||
}
|
||||
|
||||
func NewDialer(opts ...dialer.Option) dialer.Dialer {
|
||||
options := dialer.Options{}
|
||||
for _, opt := range opts {
|
||||
opt(&options)
|
||||
}
|
||||
|
||||
return &mwsDialer{
|
||||
sessions: make(map[string]*muxSession),
|
||||
options: options,
|
||||
}
|
||||
}
|
||||
|
||||
func NewTLSDialer(opts ...dialer.Option) dialer.Dialer {
|
||||
options := dialer.Options{}
|
||||
for _, opt := range opts {
|
||||
opt(&options)
|
||||
}
|
||||
|
||||
return &mwsDialer{
|
||||
tlsEnabled: true,
|
||||
sessions: make(map[string]*muxSession),
|
||||
options: options,
|
||||
}
|
||||
}
|
||||
func (d *mwsDialer) Init(md md.Metadata) (err error) {
|
||||
if err = d.parseMetadata(md); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Multiplex implements dialer.Multiplexer interface.
|
||||
func (d *mwsDialer) Multiplex() bool {
|
||||
return true
|
||||
}
|
||||
|
||||
func (d *mwsDialer) Dial(ctx context.Context, addr string, opts ...dialer.DialOption) (conn net.Conn, err error) {
|
||||
d.sessionMutex.Lock()
|
||||
defer d.sessionMutex.Unlock()
|
||||
|
||||
session, ok := d.sessions[addr]
|
||||
if session != nil && session.IsClosed() {
|
||||
delete(d.sessions, addr) // session is dead
|
||||
ok = false
|
||||
}
|
||||
if !ok {
|
||||
var options dialer.DialOptions
|
||||
for _, opt := range opts {
|
||||
opt(&options)
|
||||
}
|
||||
|
||||
conn, err = options.NetDialer.Dial(ctx, "tcp", addr)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
session = &muxSession{conn: conn}
|
||||
d.sessions[addr] = session
|
||||
}
|
||||
|
||||
return session.conn, err
|
||||
}
|
||||
|
||||
// Handshake implements dialer.Handshaker
|
||||
func (d *mwsDialer) Handshake(ctx context.Context, conn net.Conn, options ...dialer.HandshakeOption) (net.Conn, error) {
|
||||
opts := &dialer.HandshakeOptions{}
|
||||
for _, option := range options {
|
||||
option(opts)
|
||||
}
|
||||
|
||||
d.sessionMutex.Lock()
|
||||
defer d.sessionMutex.Unlock()
|
||||
|
||||
session, ok := d.sessions[opts.Addr]
|
||||
if session != nil && session.conn != conn {
|
||||
conn.Close()
|
||||
return nil, errors.New("mtls: unrecognized connection")
|
||||
}
|
||||
|
||||
if !ok || session.session == nil {
|
||||
host := d.md.host
|
||||
if host == "" {
|
||||
host = opts.Addr
|
||||
}
|
||||
s, err := d.initSession(ctx, host, conn)
|
||||
if err != nil {
|
||||
d.options.Logger.Error(err)
|
||||
conn.Close()
|
||||
delete(d.sessions, opts.Addr)
|
||||
return nil, err
|
||||
}
|
||||
session = s
|
||||
d.sessions[opts.Addr] = session
|
||||
}
|
||||
cc, err := session.GetConn()
|
||||
if err != nil {
|
||||
session.Close()
|
||||
delete(d.sessions, opts.Addr)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return cc, nil
|
||||
}
|
||||
|
||||
func (d *mwsDialer) initSession(ctx context.Context, host string, conn net.Conn) (*muxSession, error) {
|
||||
dialer := websocket.Dialer{
|
||||
HandshakeTimeout: d.md.handshakeTimeout,
|
||||
ReadBufferSize: d.md.readBufferSize,
|
||||
WriteBufferSize: d.md.writeBufferSize,
|
||||
EnableCompression: d.md.enableCompression,
|
||||
NetDial: func(net, addr string) (net.Conn, error) {
|
||||
return conn, nil
|
||||
},
|
||||
}
|
||||
|
||||
url := url.URL{Scheme: "ws", Host: host, Path: d.md.path}
|
||||
if d.tlsEnabled {
|
||||
url.Scheme = "wss"
|
||||
dialer.TLSClientConfig = d.options.TLSConfig
|
||||
}
|
||||
|
||||
if d.md.handshakeTimeout > 0 {
|
||||
conn.SetReadDeadline(time.Now().Add(d.md.handshakeTimeout))
|
||||
}
|
||||
|
||||
c, resp, err := dialer.DialContext(ctx, url.String(), d.md.header)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resp.Body.Close()
|
||||
|
||||
if d.md.handshakeTimeout > 0 {
|
||||
conn.SetReadDeadline(time.Time{})
|
||||
}
|
||||
|
||||
cc := ws_util.Conn(c)
|
||||
|
||||
if d.md.keepAlive > 0 {
|
||||
c.SetReadDeadline(time.Now().Add(d.md.keepAlive * 2))
|
||||
c.SetPongHandler(func(string) error {
|
||||
c.SetReadDeadline(time.Now().Add(d.md.keepAlive * 2))
|
||||
return nil
|
||||
})
|
||||
go d.keepAlive(cc)
|
||||
}
|
||||
|
||||
// stream multiplex
|
||||
smuxConfig := smux.DefaultConfig()
|
||||
smuxConfig.KeepAliveDisabled = d.md.muxKeepAliveDisabled
|
||||
if d.md.muxKeepAliveInterval > 0 {
|
||||
smuxConfig.KeepAliveInterval = d.md.muxKeepAliveInterval
|
||||
}
|
||||
if d.md.muxKeepAliveTimeout > 0 {
|
||||
smuxConfig.KeepAliveTimeout = d.md.muxKeepAliveTimeout
|
||||
}
|
||||
if d.md.muxMaxFrameSize > 0 {
|
||||
smuxConfig.MaxFrameSize = d.md.muxMaxFrameSize
|
||||
}
|
||||
if d.md.muxMaxReceiveBuffer > 0 {
|
||||
smuxConfig.MaxReceiveBuffer = d.md.muxMaxReceiveBuffer
|
||||
}
|
||||
if d.md.muxMaxStreamBuffer > 0 {
|
||||
smuxConfig.MaxStreamBuffer = d.md.muxMaxStreamBuffer
|
||||
}
|
||||
|
||||
session, err := smux.Client(cc, smuxConfig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &muxSession{conn: cc, session: session}, nil
|
||||
}
|
||||
|
||||
func (d *mwsDialer) keepAlive(conn ws_util.WebsocketConn) {
|
||||
ticker := time.NewTicker(d.md.keepAlive)
|
||||
defer ticker.Stop()
|
||||
|
||||
for range ticker.C {
|
||||
conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
|
||||
if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
87
dialer/mws/metadata.go
Normal file
87
dialer/mws/metadata.go
Normal file
@ -0,0 +1,87 @@
|
||||
package mws
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
mdata "github.com/go-gost/gost/v3/pkg/metadata"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultPath = "/ws"
|
||||
)
|
||||
|
||||
type metadata struct {
|
||||
host string
|
||||
path string
|
||||
|
||||
handshakeTimeout time.Duration
|
||||
readHeaderTimeout time.Duration
|
||||
readBufferSize int
|
||||
writeBufferSize int
|
||||
enableCompression bool
|
||||
|
||||
muxKeepAliveDisabled bool
|
||||
muxKeepAliveInterval time.Duration
|
||||
muxKeepAliveTimeout time.Duration
|
||||
muxMaxFrameSize int
|
||||
muxMaxReceiveBuffer int
|
||||
muxMaxStreamBuffer int
|
||||
|
||||
header http.Header
|
||||
keepAlive time.Duration
|
||||
}
|
||||
|
||||
func (d *mwsDialer) parseMetadata(md mdata.Metadata) (err error) {
|
||||
const (
|
||||
host = "host"
|
||||
path = "path"
|
||||
|
||||
handshakeTimeout = "handshakeTimeout"
|
||||
readHeaderTimeout = "readHeaderTimeout"
|
||||
readBufferSize = "readBufferSize"
|
||||
writeBufferSize = "writeBufferSize"
|
||||
enableCompression = "enableCompression"
|
||||
|
||||
header = "header"
|
||||
keepAlive = "keepAlive"
|
||||
|
||||
muxKeepAliveDisabled = "muxKeepAliveDisabled"
|
||||
muxKeepAliveInterval = "muxKeepAliveInterval"
|
||||
muxKeepAliveTimeout = "muxKeepAliveTimeout"
|
||||
muxMaxFrameSize = "muxMaxFrameSize"
|
||||
muxMaxReceiveBuffer = "muxMaxReceiveBuffer"
|
||||
muxMaxStreamBuffer = "muxMaxStreamBuffer"
|
||||
)
|
||||
|
||||
d.md.host = mdata.GetString(md, host)
|
||||
|
||||
d.md.path = mdata.GetString(md, path)
|
||||
if d.md.path == "" {
|
||||
d.md.path = defaultPath
|
||||
}
|
||||
|
||||
d.md.muxKeepAliveDisabled = mdata.GetBool(md, muxKeepAliveDisabled)
|
||||
d.md.muxKeepAliveInterval = mdata.GetDuration(md, muxKeepAliveInterval)
|
||||
d.md.muxKeepAliveTimeout = mdata.GetDuration(md, muxKeepAliveTimeout)
|
||||
d.md.muxMaxFrameSize = mdata.GetInt(md, muxMaxFrameSize)
|
||||
d.md.muxMaxReceiveBuffer = mdata.GetInt(md, muxMaxReceiveBuffer)
|
||||
d.md.muxMaxStreamBuffer = mdata.GetInt(md, muxMaxStreamBuffer)
|
||||
|
||||
d.md.handshakeTimeout = mdata.GetDuration(md, handshakeTimeout)
|
||||
d.md.readHeaderTimeout = mdata.GetDuration(md, readHeaderTimeout)
|
||||
d.md.readBufferSize = mdata.GetInt(md, readBufferSize)
|
||||
d.md.writeBufferSize = mdata.GetInt(md, writeBufferSize)
|
||||
d.md.enableCompression = mdata.GetBool(md, enableCompression)
|
||||
|
||||
if m := mdata.GetStringMapString(md, header); len(m) > 0 {
|
||||
h := http.Header{}
|
||||
for k, v := range m {
|
||||
h.Add(k, v)
|
||||
}
|
||||
d.md.header = h
|
||||
}
|
||||
d.md.keepAlive = mdata.GetDuration(md, keepAlive)
|
||||
|
||||
return
|
||||
}
|
||||
Reference in New Issue
Block a user