add mws dialer

This commit is contained in:
ginuerzh
2021-12-17 14:05:59 +08:00
parent bfe5eae172
commit 8e31e532e4
23 changed files with 795 additions and 137 deletions

109
pkg/dialer/ws/dialer.go Normal file
View File

@ -0,0 +1,109 @@
package ws
import (
"context"
"net"
"net/url"
"time"
"github.com/go-gost/gost/pkg/dialer"
ws_util "github.com/go-gost/gost/pkg/internal/util/ws"
"github.com/go-gost/gost/pkg/logger"
md "github.com/go-gost/gost/pkg/metadata"
"github.com/go-gost/gost/pkg/registry"
"github.com/gorilla/websocket"
)
func init() {
registry.RegisterDialer("ws", NewDialer)
registry.RegisterDialer("wss", NewTLSDialer)
}
type wsDialer struct {
tlsEnabled bool
logger logger.Logger
md metadata
}
func NewDialer(opts ...dialer.Option) dialer.Dialer {
options := &dialer.Options{}
for _, opt := range opts {
opt(options)
}
return &wsDialer{
logger: options.Logger,
}
}
func NewTLSDialer(opts ...dialer.Option) dialer.Dialer {
options := &dialer.Options{}
for _, opt := range opts {
opt(options)
}
return &wsDialer{
tlsEnabled: true,
logger: options.Logger,
}
}
func (d *wsDialer) Init(md md.Metadata) (err error) {
return d.parseMetadata(md)
}
func (d *wsDialer) Dial(ctx context.Context, addr string, opts ...dialer.DialOption) (net.Conn, error) {
var options dialer.DialOptions
for _, opt := range opts {
opt(&options)
}
var netd net.Dialer
conn, err := netd.DialContext(ctx, "tcp", addr)
if err != nil {
d.logger.Error(err)
}
return conn, err
}
// Handshake implements dialer.Handshaker
func (d *wsDialer) Handshake(ctx context.Context, conn net.Conn, options ...dialer.HandshakeOption) (net.Conn, error) {
opts := &dialer.HandshakeOptions{}
for _, option := range options {
option(opts)
}
if d.md.handshakeTimeout > 0 {
conn.SetDeadline(time.Now().Add(d.md.handshakeTimeout))
defer conn.SetDeadline(time.Time{})
}
host := d.md.host
if host == "" {
host = opts.Addr
}
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.md.tlsConfig
}
c, resp, err := dialer.Dial(url.String(), d.md.header)
if err != nil {
return nil, err
}
resp.Body.Close()
return ws_util.Conn(c), nil
}

86
pkg/dialer/ws/metadata.go Normal file
View File

@ -0,0 +1,86 @@
package ws
import (
"crypto/tls"
"fmt"
"net"
"net/http"
"time"
tls_util "github.com/go-gost/gost/pkg/common/util/tls"
md "github.com/go-gost/gost/pkg/metadata"
)
const (
defaultPath = "/ws"
)
type metadata struct {
path string
host string
tlsConfig *tls.Config
handshakeTimeout time.Duration
readHeaderTimeout time.Duration
readBufferSize int
writeBufferSize int
enableCompression bool
header http.Header
}
func (d *wsDialer) parseMetadata(md md.Metadata) (err error) {
const (
path = "path"
host = "host"
certFile = "certFile"
keyFile = "keyFile"
caFile = "caFile"
secure = "secure"
serverName = "serverName"
handshakeTimeout = "handshakeTimeout"
readHeaderTimeout = "readHeaderTimeout"
readBufferSize = "readBufferSize"
writeBufferSize = "writeBufferSize"
enableCompression = "enableCompression"
header = "header"
)
d.md.path = md.GetString(path)
if d.md.path == "" {
d.md.path = defaultPath
}
d.md.host = md.GetString(host)
sn, _, _ := net.SplitHostPort(md.GetString(serverName))
if sn == "" {
sn = "localhost"
}
d.md.tlsConfig, err = tls_util.LoadClientConfig(
md.GetString(certFile),
md.GetString(keyFile),
md.GetString(caFile),
md.GetBool(secure),
sn,
)
d.md.handshakeTimeout = md.GetDuration(handshakeTimeout)
d.md.readHeaderTimeout = md.GetDuration(readHeaderTimeout)
d.md.readBufferSize = md.GetInt(readBufferSize)
d.md.writeBufferSize = md.GetInt(writeBufferSize)
d.md.enableCompression = md.GetBool(enableCompression)
if mm, _ := md.Get(header).(map[interface{}]interface{}); len(mm) > 0 {
h := http.Header{}
for k, v := range mm {
h.Add(fmt.Sprintf("%v", k), fmt.Sprintf("%v", v))
}
d.md.header = h
}
return
}

38
pkg/dialer/ws/mux/conn.go Normal file
View File

@ -0,0 +1,38 @@
package mux
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()
}

220
pkg/dialer/ws/mux/dialer.go Normal file
View File

@ -0,0 +1,220 @@
package mux
import (
"context"
"errors"
"net"
"net/url"
"sync"
"time"
"github.com/go-gost/gost/pkg/dialer"
ws_util "github.com/go-gost/gost/pkg/internal/util/ws"
"github.com/go-gost/gost/pkg/logger"
md "github.com/go-gost/gost/pkg/metadata"
"github.com/go-gost/gost/pkg/registry"
"github.com/gorilla/websocket"
"github.com/xtaci/smux"
)
func init() {
registry.RegisterDialer("mws", NewDialer)
registry.RegisterDialer("mwss", NewTLSDialer)
}
type mwsDialer struct {
sessions map[string]*muxSession
sessionMutex sync.Mutex
logger logger.Logger
md metadata
tlsEnabled bool
}
func NewDialer(opts ...dialer.Option) dialer.Dialer {
options := &dialer.Options{}
for _, opt := range opts {
opt(options)
}
return &mwsDialer{
sessions: make(map[string]*muxSession),
logger: options.Logger,
}
}
func NewTLSDialer(opts ...dialer.Option) dialer.Dialer {
options := &dialer.Options{}
for _, opt := range opts {
opt(options)
}
return &mwsDialer{
sessions: make(map[string]*muxSession),
logger: options.Logger,
tlsEnabled: true,
}
}
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) {
var options dialer.DialOptions
for _, opt := range opts {
opt(&options)
}
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 {
conn, err = d.dial(ctx, "tcp", addr, &options)
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()
if d.md.handshakeTimeout > 0 {
conn.SetDeadline(time.Now().Add(d.md.handshakeTimeout))
defer conn.SetDeadline(time.Time{})
}
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.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) dial(ctx context.Context, network, addr string, opts *dialer.DialOptions) (net.Conn, error) {
dial := opts.DialFunc
if dial != nil {
conn, err := dial(ctx, addr)
if err != nil {
d.logger.Error(err)
} else {
d.logger.WithFields(map[string]interface{}{
"src": conn.LocalAddr().String(),
"dst": addr,
}).Debug("dial with dial func")
}
return conn, err
}
var netd net.Dialer
conn, err := netd.DialContext(ctx, network, addr)
if err != nil {
d.logger.Error(err)
} else {
d.logger.WithFields(map[string]interface{}{
"src": conn.LocalAddr().String(),
"dst": addr,
}).Debugf("dial direct %s/%s", addr, network)
}
return conn, err
}
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.md.tlsConfig
}
c, resp, err := dialer.Dial(url.String(), d.md.header)
if err != nil {
return nil, err
}
resp.Body.Close()
conn = ws_util.Conn(c)
// 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(conn, smuxConfig)
if err != nil {
return nil, err
}
return &muxSession{conn: conn, session: session}, nil
}

View File

@ -0,0 +1,106 @@
package mux
import (
"crypto/tls"
"fmt"
"net"
"net/http"
"time"
tls_util "github.com/go-gost/gost/pkg/common/util/tls"
md "github.com/go-gost/gost/pkg/metadata"
)
const (
defaultPath = "/ws"
)
type metadata struct {
path string
host string
tlsConfig *tls.Config
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
}
func (d *mwsDialer) parseMetadata(md md.Metadata) (err error) {
const (
path = "path"
host = "host"
certFile = "certFile"
keyFile = "keyFile"
caFile = "caFile"
secure = "secure"
serverName = "serverName"
handshakeTimeout = "handshakeTimeout"
readHeaderTimeout = "readHeaderTimeout"
readBufferSize = "readBufferSize"
writeBufferSize = "writeBufferSize"
enableCompression = "enableCompression"
header = "header"
muxKeepAliveDisabled = "muxKeepAliveDisabled"
muxKeepAliveInterval = "muxKeepAliveInterval"
muxKeepAliveTimeout = "muxKeepAliveTimeout"
muxMaxFrameSize = "muxMaxFrameSize"
muxMaxReceiveBuffer = "muxMaxReceiveBuffer"
muxMaxStreamBuffer = "muxMaxStreamBuffer"
)
d.md.path = md.GetString(path)
if d.md.path == "" {
d.md.path = defaultPath
}
d.md.host = md.GetString(host)
sn, _, _ := net.SplitHostPort(md.GetString(serverName))
if sn == "" {
sn = "localhost"
}
d.md.tlsConfig, err = tls_util.LoadClientConfig(
md.GetString(certFile),
md.GetString(keyFile),
md.GetString(caFile),
md.GetBool(secure),
sn,
)
d.md.muxKeepAliveDisabled = md.GetBool(muxKeepAliveDisabled)
d.md.muxKeepAliveInterval = md.GetDuration(muxKeepAliveInterval)
d.md.muxKeepAliveTimeout = md.GetDuration(muxKeepAliveTimeout)
d.md.muxMaxFrameSize = md.GetInt(muxMaxFrameSize)
d.md.muxMaxReceiveBuffer = md.GetInt(muxMaxReceiveBuffer)
d.md.muxMaxStreamBuffer = md.GetInt(muxMaxStreamBuffer)
d.md.handshakeTimeout = md.GetDuration(handshakeTimeout)
d.md.readHeaderTimeout = md.GetDuration(readHeaderTimeout)
d.md.readBufferSize = md.GetInt(readBufferSize)
d.md.writeBufferSize = md.GetInt(writeBufferSize)
d.md.enableCompression = md.GetBool(enableCompression)
if mm, _ := md.Get(header).(map[interface{}]interface{}); len(mm) > 0 {
h := http.Header{}
for k, v := range mm {
h.Add(fmt.Sprintf("%v", k), fmt.Sprintf("%v", v))
}
d.md.header = h
}
return
}