bug

lib/common/const.go

@@ -6,20 +6,21 @@ const (
 	COMPRESS_NONE_DECODE
 	COMPRESS_SNAPY_ENCODE
 	COMPRESS_SNAPY_DECODE
-	VERIFY_EER     = "vkey"
-	VERIFY_SUCCESS = "sucs"
-	WORK_MAIN      = "main"
-	WORK_CHAN      = "chan"
-	WORK_CONFIG    = "conf"
-	WORK_REGISTER  = "rgst"
-	WORK_STATUS    = "stus"
-	RES_SIGN       = "sign"
-	RES_MSG        = "msg0"
-	RES_CLOSE      = "clse"
-	NEW_CONN       = "conn" //新连接标志
-	NEW_TASK       = "task" //新连接标志
-	NEW_CONF       = "conf" //新连接标志
-	NEW_HOST       = "host" //新连接标志
+	VERIFY_EER       = "vkey"
+	VERIFY_SUCCESS   = "sucs"
+	WORK_MAIN        = "main"
+	WORK_CHAN        = "chan"
+	WORK_SEND_STATUS = "sdst"
+	WORK_CONFIG      = "conf"
+	WORK_REGISTER    = "rgst"
+	WORK_STATUS      = "stus"
+	RES_SIGN         = "sign"
+	RES_MSG          = "msg0"
+	RES_CLOSE        = "clse"
+	NEW_CONN         = "conn" //新连接标志
+	NEW_TASK         = "task" //新连接标志
+	NEW_CONF         = "conf" //新连接标志
+	NEW_HOST         = "host" //新连接标志
 
 	CONN_TCP          = "tcp"
 	CONN_UDP          = "udp"

lib/conn/conn.go

@@ -74,12 +74,12 @@ func (s *Conn) ReadLen(cLen int) ([]byte, error) {
 		return nil, errors.New("长度错误" + strconv.Itoa(cLen))
 	}
 	var buf []byte
-	if cLen <= pool.PoolSizeSmall {
+	if cLen < pool.PoolSizeSmall {
 		buf = pool.BufPoolSmall.Get().([]byte)[:cLen]
-		defer pool.BufPoolSmall.Put(buf)
+		defer pool.PutBufPoolSmall(buf)
 	} else {
 		buf = pool.BufPoolMax.Get().([]byte)[:cLen]
-		defer pool.BufPoolMax.Put(buf)
+		defer pool.PutBufPoolMax(buf)
 	}
 	if n, err := io.ReadFull(s, buf); err != nil || n != cLen {
 		return buf, errors.New("Error reading specified length " + err.Error())
@@ -190,14 +190,10 @@ func (s *Conn) SendMsg(content []byte, link *Link) (n int, err error) {
 	*/
 	s.Lock()
 	defer s.Unlock()
-	raw := bytes.NewBuffer([]byte{})
-	binary.Write(raw, binary.LittleEndian, int32(link.Id))
-	if n, err = s.Write(raw.Bytes()); err != nil {
+	if err = binary.Write(s.Conn, binary.LittleEndian, int32(link.Id)); err != nil {
 		return
 	}
-	raw.Reset()
-	binary.Write(raw, binary.LittleEndian, content)
-	n, err = s.WriteTo(raw.Bytes(), link.En, link.Crypt, link.Rate)
+	n, err = s.WriteTo(content, link.En, link.Crypt, link.Rate)
 	return
 }
 
@@ -260,6 +256,8 @@ func (s *Conn) GetLinkInfo() (lk *Link, err error) {
 		lk.En = common.GetIntNoErrByStr(string(buf[11+hostLen]))
 		lk.De = common.GetIntNoErrByStr(string(buf[12+hostLen]))
 		lk.Crypt = common.GetBoolByStr(string(buf[13+hostLen]))
+		lk.MsgCh = make(chan []byte)
+		lk.StatusCh = make(chan bool)
 	}
 	return
 }
@@ -399,6 +397,10 @@ func (s *Conn) GetTaskInfo() (t *file.Tunnel, err error) {
 	return
 }
 
+func (s *Conn) WriteWriteSuccess(id int) error {
+	return binary.Write(s.Conn, binary.LittleEndian, int32(id))
+}
+
 //write connect success
 func (s *Conn) WriteSuccess(id int) (int, error) {
 	raw := bytes.NewBuffer([]byte{})

lib/conn/link.go

@@ -1,7 +1,9 @@
 package conn
 
 import (
+	"github.com/cnlh/nps/lib/common"
 	"github.com/cnlh/nps/lib/file"
+	"github.com/cnlh/nps/lib/pool"
 	"github.com/cnlh/nps/lib/rate"
 	"net"
 )
@@ -18,6 +20,9 @@ type Link struct {
 	UdpListener   *net.UDPConn
 	Rate          *rate.Rate
 	UdpRemoteAddr *net.UDPAddr
+	MsgCh         chan []byte
+	MsgConn       *Conn
+	StatusCh      chan bool
 }
 
 func NewLink(id int, connType string, host string, en, de int, crypt bool, c *Conn, flow *file.Flow, udpListener *net.UDPConn, rate *rate.Rate, UdpRemoteAddr *net.UDPAddr) *Link {
@@ -33,5 +38,37 @@ func NewLink(id int, connType string, host string, en, de int, crypt bool, c *Co
 		UdpListener:   udpListener,
 		Rate:          rate,
 		UdpRemoteAddr: UdpRemoteAddr,
+		MsgCh:         make(chan []byte),
+		StatusCh:      make(chan bool),
 	}
 }
+
+func (s *Link) Run(flow bool) {
+	go func() {
+		for {
+			select {
+			case content := <-s.MsgCh:
+				if len(content) == len(common.IO_EOF) && string(content) == common.IO_EOF {
+					if s.Conn != nil {
+						s.Conn.Close()
+					}
+					return
+				} else {
+					if s.UdpListener != nil && s.UdpRemoteAddr != nil {
+						s.UdpListener.WriteToUDP(content, s.UdpRemoteAddr)
+					} else {
+						s.Conn.Write(content)
+					}
+					if flow {
+						s.Flow.Add(0, len(content))
+					}
+					if s.ConnType == common.CONN_UDP {
+						return
+					}
+					s.MsgConn.WriteWriteSuccess(s.Id)
+					pool.PutBufPoolCopy(content)
+				}
+			}
+		}
+	}()
+}

lib/pool/pool.go

@@ -36,14 +36,26 @@ var BufPoolCopy = sync.Pool{
 	},
 }
 
+func PutBufPoolUdp(buf []byte) {
+	if cap(buf) == PoolSizeUdp {
+		BufPoolUdp.Put(buf[:PoolSizeUdp])
+	}
+}
+
 func PutBufPoolCopy(buf []byte) {
 	if cap(buf) == PoolSizeCopy {
 		BufPoolCopy.Put(buf[:PoolSizeCopy])
 	}
 }
 
-func PutBufPoolUdp(buf []byte) {
-	if cap(buf) == PoolSizeUdp {
-		BufPoolUdp.Put(buf[:PoolSizeUdp])
+func PutBufPoolSmall(buf []byte) {
+	if cap(buf) == PoolSizeSmall {
+		BufPoolSmall.Put(buf[:PoolSizeSmall])
+	}
+}
+
+func PutBufPoolMax(buf []byte) {
+	if cap(buf) == PoolSize {
+		BufPoolMax.Put(buf[:PoolSize])
 	}
 }