client net dail timeout

pull/125/head
lpxxn 2019-05-22 17:40:45 +08:00
parent 4b7b2f4c27
commit b56e8688e3
3 changed files with 53 additions and 14 deletions

View File

@ -2,6 +2,11 @@ package client
import ( import (
"bufio" "bufio"
"net"
"net/http"
"strconv"
"time"
"github.com/cnlh/nps/lib/common" "github.com/cnlh/nps/lib/common"
"github.com/cnlh/nps/lib/config" "github.com/cnlh/nps/lib/config"
"github.com/cnlh/nps/lib/conn" "github.com/cnlh/nps/lib/conn"
@ -9,10 +14,6 @@ import (
"github.com/cnlh/nps/lib/mux" "github.com/cnlh/nps/lib/mux"
"github.com/cnlh/nps/vender/github.com/astaxie/beego/logs" "github.com/cnlh/nps/vender/github.com/astaxie/beego/logs"
"github.com/cnlh/nps/vender/github.com/xtaci/kcp" "github.com/cnlh/nps/vender/github.com/xtaci/kcp"
"net"
"net/http"
"strconv"
"time"
) )
type TRPClient struct { type TRPClient struct {
@ -159,7 +160,7 @@ func (s *TRPClient) handleChan(src net.Conn) {
lk.Host = common.FormatAddress(lk.Host) lk.Host = common.FormatAddress(lk.Host)
//if Conn type is http, read the request and log //if Conn type is http, read the request and log
if lk.ConnType == "http" { 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()) logs.Warn("connect to %s error %s", lk.Host, err.Error())
src.Close() src.Close()
} else { } else {
@ -183,7 +184,7 @@ func (s *TRPClient) handleChan(src net.Conn) {
return return
} }
//connect to target if conn type is tcp or udp //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()) logs.Warn("connect to %s error %s", lk.Host, err.Error())
src.Close() src.Close()
} else { } else {

View File

@ -1,5 +1,7 @@
package conn package conn
import "time"
type Secret struct { type Secret struct {
Password string Password string
Conn *Conn Conn *Conn
@ -19,9 +21,20 @@ type Link struct {
Compress bool Compress bool
LocalProxy bool LocalProxy bool
RemoteAddr string 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{ return &Link{
RemoteAddr: remoteAddr, RemoteAddr: remoteAddr,
ConnType: connType, ConnType: connType,
@ -29,5 +42,22 @@ func NewLink(connType string, host string, crypt bool, compress bool, remoteAddr
Crypt: crypt, Crypt: crypt,
Compress: compress, Compress: compress,
LocalProxy: localProxy, 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
} }
} }

View File

@ -93,7 +93,8 @@ func (s *conn) Write(buf []byte) (int, error) {
if s.isClose { if s.isClose {
return 0, errors.New("the conn has closed") return 0, errors.New("the conn has closed")
} }
ch := make(chan struct{}) ch := make(chan error)
var err error
go s.write(buf, ch) go s.write(buf, ch)
if t := s.writeTimeOut.Sub(time.Now()); t > 0 { if t := s.writeTimeOut.Sub(time.Now()); t > 0 {
timer := time.NewTimer(t) timer := time.NewTimer(t)
@ -101,17 +102,20 @@ func (s *conn) Write(buf []byte) (int, error) {
select { select {
case <-timer.C: case <-timer.C:
return 0, errors.New("write timeout") return 0, errors.New("write timeout")
case <-ch: case err = <-ch:
} }
} else { } else {
<-ch err = <-ch
} }
if s.isClose { if s.isClose {
return 0, io.EOF return 0, io.EOF
} }
if err != nil {
return 0, err
}
return len(buf), nil return len(buf), nil
} }
func (s *conn) write(buf []byte, ch chan struct{}) { func (s *conn) write(buf []byte, ch chan error) {
start := 0 start := 0
l := len(buf) l := len(buf)
for { for {
@ -120,14 +124,18 @@ func (s *conn) write(buf []byte, ch chan struct{}) {
} }
s.hasWrite++ s.hasWrite++
if l-start > pool.PoolSizeCopy { if l-start > pool.PoolSizeCopy {
s.mux.sendInfo(MUX_NEW_MSG, s.connId, buf[start:start+pool.PoolSizeCopy]) if err := s.mux.sendInfo(MUX_NEW_MSG, s.connId, buf[start:start+pool.PoolSizeCopy]); err != nil {
ch <- err
}
start += pool.PoolSizeCopy start += pool.PoolSizeCopy
} else { } else {
s.mux.sendInfo(MUX_NEW_MSG, s.connId, buf[start:l]) if err := s.mux.sendInfo(MUX_NEW_MSG, s.connId, buf[start:l]); err != nil {
ch <- err
}
break break
} }
} }
ch <- struct{}{} ch <- nil
} }
func (s *conn) Close() error { func (s *conn) Close() error {