core/chain/transport.go
2022-09-08 22:39:17 +08:00

125 lines
2.8 KiB
Go

package chain
import (
"context"
"net"
"time"
net_dialer "github.com/go-gost/core/common/net/dialer"
"github.com/go-gost/core/connector"
"github.com/go-gost/core/dialer"
)
type Transport struct {
addr string
ifceName string
sockOpts *SockOpts
route Route
dialer dialer.Dialer
connector connector.Connector
timeout time.Duration
}
func (tr *Transport) Copy() *Transport {
tr2 := &Transport{}
*tr2 = *tr
return tr
}
func (tr *Transport) WithInterface(ifceName string) *Transport {
tr.ifceName = ifceName
return tr
}
func (tr *Transport) WithSockOpts(so *SockOpts) *Transport {
tr.sockOpts = so
return tr
}
func (tr *Transport) WithDialer(dialer dialer.Dialer) *Transport {
tr.dialer = dialer
return tr
}
func (tr *Transport) WithConnector(connector connector.Connector) *Transport {
tr.connector = connector
return tr
}
func (tr *Transport) WithTimeout(d time.Duration) *Transport {
tr.timeout = d
return tr
}
func (tr *Transport) Dial(ctx context.Context, addr string) (net.Conn, error) {
netd := &net_dialer.NetDialer{
Interface: tr.ifceName,
Timeout: tr.timeout,
}
if tr.sockOpts != nil {
netd.Mark = tr.sockOpts.Mark
}
if tr.route != nil && tr.route.Len() > 0 {
netd.DialFunc = func(ctx context.Context, network, addr string) (net.Conn, error) {
return tr.route.Dial(ctx, network, addr)
}
}
opts := []dialer.DialOption{
dialer.HostDialOption(tr.addr),
dialer.NetDialerDialOption(netd),
}
return tr.dialer.Dial(ctx, addr, opts...)
}
func (tr *Transport) Handshake(ctx context.Context, conn net.Conn) (net.Conn, error) {
var err error
if hs, ok := tr.dialer.(dialer.Handshaker); ok {
conn, err = hs.Handshake(ctx, conn,
dialer.AddrHandshakeOption(tr.addr))
if err != nil {
return nil, err
}
}
if hs, ok := tr.connector.(connector.Handshaker); ok {
return hs.Handshake(ctx, conn)
}
return conn, nil
}
func (tr *Transport) Connect(ctx context.Context, conn net.Conn, network, address string) (net.Conn, error) {
netd := &net_dialer.NetDialer{
Interface: tr.ifceName,
Timeout: tr.timeout,
}
if tr.sockOpts != nil {
netd.Mark = tr.sockOpts.Mark
}
return tr.connector.Connect(ctx, conn, network, address,
connector.NetDialerConnectOption(netd),
)
}
func (tr *Transport) Bind(ctx context.Context, conn net.Conn, network, address string, opts ...connector.BindOption) (net.Listener, error) {
if binder, ok := tr.connector.(connector.Binder); ok {
return binder.Bind(ctx, conn, network, address, opts...)
}
return nil, connector.ErrBindUnsupported
}
func (tr *Transport) Multiplex() bool {
if mux, ok := tr.dialer.(dialer.Multiplexer); ok {
return mux.Multiplex()
}
return false
}
func (tr *Transport) WithRoute(r Route) *Transport {
tr.route = r
return tr
}
func (tr *Transport) WithAddr(addr string) *Transport {
tr.addr = addr
return tr
}