diff --git a/client/client.go b/client/client.go index 5bd0b01..030d35f 100755 --- a/client/client.go +++ b/client/client.go @@ -8,12 +8,13 @@ import ( "time" "github.com/astaxie/beego/logs" + "github.com/xtaci/kcp-go" + "github.com/cnlh/nps/lib/common" "github.com/cnlh/nps/lib/config" "github.com/cnlh/nps/lib/conn" "github.com/cnlh/nps/lib/crypt" "github.com/cnlh/nps/lib/mux" - "github.com/xtaci/kcp-go" ) type TRPClient struct { @@ -165,7 +166,7 @@ func (s *TRPClient) handleChan(src net.Conn) { lk.Host = common.FormatAddress(lk.Host) //if Conn type is http, read the request and log if lk.ConnType == "http" { - if targetConn, err := net.Dial(common.CONN_TCP, lk.Host); err != nil { + if targetConn, err := net.DialTimeout(common.CONN_TCP, lk.Host, lk.Option.Timeout); err != nil { logs.Warn("connect to %s error %s", lk.Host, err.Error()) src.Close() } else { @@ -189,7 +190,7 @@ func (s *TRPClient) handleChan(src net.Conn) { return } //connect to target if conn type is tcp or udp - if targetConn, err := net.Dial(lk.ConnType, lk.Host); err != nil { + if targetConn, err := net.DialTimeout(lk.ConnType, lk.Host, lk.Option.Timeout); err != nil { logs.Warn("connect to %s error %s", lk.Host, err.Error()) src.Close() } else { diff --git a/lib/conn/link.go b/lib/conn/link.go index 974827d..653894e 100644 --- a/lib/conn/link.go +++ b/lib/conn/link.go @@ -1,5 +1,7 @@ package conn +import "time" + type Secret struct { Password string Conn *Conn @@ -19,9 +21,20 @@ type Link struct { Compress bool LocalProxy bool RemoteAddr string + Option Options } -func NewLink(connType string, host string, crypt bool, compress bool, remoteAddr string, localProxy bool) *Link { +type Option func(*Options) + +type Options struct { + Timeout time.Duration +} + +var defaultTimeOut = time.Second * 5 + +func NewLink(connType string, host string, crypt bool, compress bool, remoteAddr string, localProxy bool, opts ...Option) *Link { + options := newOptions(opts...) + return &Link{ RemoteAddr: remoteAddr, ConnType: connType, @@ -29,5 +42,22 @@ func NewLink(connType string, host string, crypt bool, compress bool, remoteAddr Crypt: crypt, Compress: compress, LocalProxy: localProxy, + Option: options, + } +} + +func newOptions(opts ...Option) Options { + opt := Options{ + Timeout: defaultTimeOut, + } + for _, o := range opts { + o(&opt) + } + return opt +} + +func LinkTimeout(t time.Duration) Option { + return func(opt *Options) { + opt.Timeout = t } }