重构，用多态模式写数据发送和接收的代码。速度快了，可能和减少了判断有关。

benckmark.go

@@ -68,7 +68,7 @@ func attack() {
 	log.Println("Finish")
 }
 
-func main() {
+func main1() {
 	cfg := profile.Config{
 		MemProfile:     true,
 		ProfilePath:    "./profile", // store profiles in current directory

communicator/receiver.go

@@ -0,0 +1,104 @@
+package communicator
+
+import (
+	"io"
+
+	"github.com/TransX/log"
+	"github.com/TransX/tscipher"
+)
+
+type Receiver interface {
+	//SendData(carrier *Carrier) (n int, err error)
+	Receive() (n int, err error)
+	Close()
+}
+
+func NewNormalReceiver(carrier *tscipher.Carrier) *NormalReceiver {
+	r := new(NormalReceiver)
+	r.carrier = carrier
+	return r
+}
+
+type NormalReceiver struct {
+	Receiver
+	carrier *tscipher.Carrier
+}
+
+func (this *NormalReceiver) Receive() (n int, err error) {
+	carrier := this.carrier
+	wrapedPackage := carrier.GetReceiveBuff() //make([]byte, 0, cap(carrier.Cache))
+	var packageData []byte
+	var _rest []byte
+	cache, _ := carrier.Cache.Get()
+	for {
+		//首先检查这个是不是完整的包，是就返回好了，免得被阻塞
+		data, rest, err := tscipher.UnwrapPackage(wrapedPackage)
+		packageData = data
+		_rest = rest
+		if err, ok := err.(*tscipher.NotPackageError); len(wrapedPackage) >= 18 && ok {
+			log.Debug("return NotPackageError %s", carrier.AttachedTunnelID)
+			return 0, err
+		}
+		if err == nil {
+			//够一个完整的包
+			capBuff := cap(carrier.GetReceiveBuff())
+			_buff := make([]byte, 0, capBuff) //释放
+			_buff = append(_buff, _rest...)
+			carrier.SetReceiveBuff(_buff)
+			break
+		}
+		//如果读到的数据不够一个完整的包
+		if len(wrapedPackage) > 0 {
+			n, err = carrier.Conn.Read(cache)
+			if err != nil {
+				log.Error("ERROR %s", err.Error())
+			}
+		} else {
+			n, err = io.ReadAtLeast(carrier.Conn, cache, 18)
+		}
+
+		if err != nil {
+			n = 0
+			return n, err
+		}
+		wrapedPackage = append(wrapedPackage, cache[:n]...)
+	}
+	decrypted, err := carrier.Cipher.Decrypt(packageData)
+	if err != nil {
+		n = 0
+		return
+	}
+	n = len(decrypted)
+	carrier.Msg.Put(decrypted, n)
+	return
+}
+
+func (this *NormalReceiver) Close() {
+	this.carrier.Conn.Close()
+}
+
+func NewDirectReceiver(carrier *tscipher.Carrier) *DirectReceiver {
+	r := new(DirectReceiver)
+	r.carrier = carrier
+	return r
+}
+
+type DirectReceiver struct {
+	Receiver
+	carrier *tscipher.Carrier
+}
+
+func (this *DirectReceiver) Receive() (n int, err error) {
+	carrier := this.carrier
+	cache, _ := carrier.Cache.Get()
+	n, err = carrier.Conn.Read(cache)
+	if err != nil {
+		n = 0
+	}
+	carrier.Msg.Put(cache, n)
+	return
+}
+
+func (this *DirectReceiver) Close() {
+	this.carrier.Conn.Close()
+}

communicator/sender.go

@@ -0,0 +1,79 @@
+package communicator
+
+import (
+	"log"
+
+	"github.com/TransX/tscipher"
+)
+
+type Sender interface {
+	//SendData(carrier *Carrier) (n int, err error)
+	Send() (n int, err error)
+	Close()
+}
+
+func NewNormalSender(carrier *tscipher.Carrier) *NormalSender {
+	r := new(NormalSender)
+	r.carrier = carrier
+	return r
+}
+
+type NormalSender struct {
+	Sender
+	carrier *tscipher.Carrier
+}
+
+func (this *NormalSender) Send() (n int, err error) {
+	carrier := this.carrier
+	msg, nByte := carrier.Msg.Get()
+	if len(msg) < nByte {
+		log.Panic("Cache of send is too small")
+	}
+	// if carrier.Cipher == nil {
+	// 	n, err = carrier.Conn.Write(msg[:nByte])
+	// 	carrier.Cache.Put(make([]byte, 1024*4), 1024*4)
+	// 	return
+	// }
+	encrypedByte, err := carrier.Cipher.Encrypt(msg[:nByte])
+	if err != nil {
+		n = 0
+		return
+	}
+	//打包
+	wraped := tscipher.WrapPackage(encrypedByte[:nByte])
+	n, err = carrier.Conn.Write(wraped)
+	carrier.Cache.Put(make([]byte, 1024*4), 1024*4)
+	return
+}
+
+func (this *NormalSender) Close() {
+	this.carrier.Conn.Close()
+}
+
+func NewDirectSender(carrier *tscipher.Carrier) *DirectSender {
+	r := new(DirectSender)
+	r.carrier = carrier
+	return r
+}
+
+type DirectSender struct {
+	Sender
+	carrier *tscipher.Carrier
+}
+
+func (this *DirectSender) Send() (n int, err error) {
+	carrier := this.carrier
+	msg, nByte := carrier.Msg.Get()
+	if len(msg) < nByte {
+		log.Panic("Cache of send is too small")
+	}
+
+	n, err = carrier.Conn.Write(msg[:nByte])
+	carrier.Cache.Put(make([]byte, 1024*4), 1024*4)
+	return
+
+}
+
+func (this *DirectSender) Close() {
+	this.carrier.Conn.Close()
+}

main.go

@@ -25,7 +25,7 @@ func tunnel() {
 
 }
 
-func main1() {
+func main() {
 	// defer profile.Start(profile.CPUProfile).Stop()
 	flag.Parse()
 	fmt.Println("Hello World!")

model/tunnel.go

@@ -3,13 +3,15 @@ package model
 import (
 
 	// "fmt"
+	"net"
+	"time"
+
 	"github.com/TransX/cache"
+	"github.com/TransX/communicator"
 	"github.com/TransX/constant"
 	"github.com/TransX/log"
 	"github.com/TransX/tscipher"
 	"github.com/spf13/viper"
-	"net"
-	"time"
 )
 
 type Tunnel struct {
@@ -95,14 +97,18 @@ func (this *Tunnel) receive(revCarrier *tscipher.Carrier) {
 	defer this.onError()
 	var n int
 	var err error
+	var receiver communicator.Receiver
+	if cipherDirection != constant.RECEIVE {
+		// revCarrier.Cipher = nil
+		// n, err = tscipher.RowReceiveData(revCarrier)
+		receiver = communicator.NewDirectReceiver(revCarrier)
+	} else {
+		receiver = communicator.NewNormalReceiver(revCarrier)
+	}
+
 	for {
 		rTimer := time.Now() //receive timer
-		if cipherDirection != constant.RECEIVE {
+		n, err = receiver.Receive()
-			revCarrier.Cipher = nil
-			n, err = tscipher.RowReceiveData(revCarrier)
-		} else {
-			n, err = tscipher.ReceiveData(revCarrier)
-		}
 		log.Info("id %s time to receive %d", id, time.Since(rTimer)/1000)
 		if err != nil {
 			log.Panic("Read panic. Tunnel id: %s. Remote Add: %s Local: %s. Err:%s", id, src.RemoteAddr().String(), src.LocalAddr().String(), err.Error())
@@ -115,13 +121,17 @@ func (this *Tunnel) send(sendCarrier *tscipher.Carrier) {
 	dest := this.dest
 	cipherDirection := this.cipherDirection
 	id := this.id
+	var sender communicator.Sender
 	defer this.onError()
 	if cipherDirection != constant.SEND {
-		sendCarrier.Cipher = nil
+		// sendCarrier.Cipher = nil
+		sender = communicator.NewDirectSender(sendCarrier)
+	} else {
+		sender = communicator.NewNormalSender(sendCarrier)
 	}
 	for {
 		sTimer := time.Now() //send timer
-		n, err := tscipher.SendData(sendCarrier)
+		n, err := sender.Send()
 		log.Info("id %s time to send %d", id, time.Since(sTimer)/1000)
 		if err != nil {
 			log.Panic("Write panic. ID: %s, Err: %s, Remote Add: %s", id, err, dest.RemoteAddr().String())

protocol/tcp.go

@@ -80,15 +80,9 @@ func (this *TransTCP) Start(listenPort, destIP, destPort string, clientOrServer
 				if clientOrServer == "client" { //加密方向
 					sendID := utils.TunnelID()
 					ntSend := model.NewTunnel(sendID, listenerConn, destConn, constant.SEND)
-					// ntSend.SetRegChan(tunMng.GetRegChan())
-					// ntSend.SetUnRegChan(tunMng.GetUnregChan())
 					receiveID := utils.TunnelID()
 					ntRev := model.NewTunnel(receiveID, destConn, listenerConn, constant.RECEIVE)
-					// ntRev.SetRegChan(tunMng.GetRegChan())
-					// ntRev.SetUnRegChan(tunMng.GetUnregChan())
 					printModelDetail(sendID, receiveID)
-					// go ntSend.Run()
-					// go ntRev.Run()
 					tunnelPair := model.NewTunnelPair(ntSend, ntRev)
 					tunnelPair.SetRegChan(tunMng.GetRegChan())
 					tunnelPair.SetUnRegChan(tunMng.GetUnregChan())
@@ -97,15 +91,9 @@ func (this *TransTCP) Start(listenPort, destIP, destPort string, clientOrServer
 				if clientOrServer == "server" {
 					receiveID := utils.TunnelID()
 					ntRev := model.NewTunnel(receiveID, listenerConn, destConn, constant.RECEIVE)
-					// ntRev.SetRegChan(tunMng.GetRegChan())
-					// ntRev.SetUnRegChan(tunMng.GetUnregChan())
 					sendID := utils.TunnelID()
 					ntSend := model.NewTunnel(sendID, destConn, listenerConn, constant.SEND)
-					// ntSend.SetRegChan(tunMng.GetRegChan())
-					// ntSend.SetUnRegChan(tunMng.GetUnregChan())
 					printModelDetail(sendID, receiveID)
-					// go ntRev.Run()
-					// go ntSend.Run()
 					tunnelPair := model.NewTunnelPair(ntSend, ntRev)
 					tunnelPair.SetRegChan(tunMng.GetRegChan())
 					tunnelPair.SetUnRegChan(tunMng.GetUnregChan())

tscipher/cipher.go

@@ -3,14 +3,13 @@ package tscipher
 import (
 	"bytes"
 	"fmt"
-	"github.com/TransX/cache"
-	"github.com/TransX/log"
-	"github.com/TransX/utils"
-	"github.com/spf13/viper"
-	"io"
 	"net"
 	"strconv"
 	"strings"
+
+	"github.com/TransX/cache"
+	"github.com/TransX/utils"
+	"github.com/spf13/viper"
 )
 
 // var StartMark = []byte("#2v!") //should be constant
@@ -141,82 +140,82 @@ func UnwrapPackage(pacakge []byte) (data []byte, rest []byte, err error) {
 
 }
 
-func SendData(carrier *Carrier) (n int, err error) {
+// func SendData(carrier *Carrier) (n int, err error) {
-	msg, nByte := carrier.Msg.Get()
+// 	msg, nByte := carrier.Msg.Get()
-	if len(msg) < nByte {
+// 	if len(msg) < nByte {
-		log.Panic("Cache of send is too small")
+// 		log.Panic("Cache of send is too small")
-	}
+// 	}
-	if carrier.Cipher == nil {
+// 	if carrier.Cipher == nil {
-		n, err = carrier.Conn.Write(msg[:nByte])
+// 		n, err = carrier.Conn.Write(msg[:nByte])
-		carrier.Cache.Put(make([]byte, 1024*4), 1024*4)
+// 		carrier.Cache.Put(make([]byte, 1024*4), 1024*4)
-		return
+// 		return
-	}
+// 	}
-	encrypedByte, err := carrier.Cipher.Encrypt(msg[:nByte])
+// 	encrypedByte, err := carrier.Cipher.Encrypt(msg[:nByte])
-	if err != nil {
+// 	if err != nil {
-		n = 0
+// 		n = 0
-		return
+// 		return
-	}
+// 	}
-	//打包
+// 	//打包
-	wraped := WrapPackage(encrypedByte[:nByte])
+// 	wraped := WrapPackage(encrypedByte[:nByte])
-	n, err = carrier.Conn.Write(wraped)
+// 	n, err = carrier.Conn.Write(wraped)
-	carrier.Cache.Put(make([]byte, 1024*4), 1024*4)
+// 	carrier.Cache.Put(make([]byte, 1024*4), 1024*4)
-	return
+// 	return
-}
+// }
-
+//
-func RowReceiveData(carrier *Carrier) (n int, err error) {
+// func RowReceiveData(carrier *Carrier) (n int, err error) {
-	cache, _ := carrier.Cache.Get()
+// 	cache, _ := carrier.Cache.Get()
-	n, err = carrier.Conn.Read(cache)
+// 	n, err = carrier.Conn.Read(cache)
-	if err != nil {
+// 	if err != nil {
-		n = 0
+// 		n = 0
-	}
+// 	}
-	carrier.Msg.Put(cache, n)
+// 	carrier.Msg.Put(cache, n)
-	return
+// 	return
-}
+// }
-
+//
-func ReceiveData(carrier *Carrier) (n int, err error) {
+// func ReceiveData(carrier *Carrier) (n int, err error) {
-	wrapedPackage := carrier.GetReceiveBuff() //make([]byte, 0, cap(carrier.Cache))
+// 	wrapedPackage := carrier.GetReceiveBuff() //make([]byte, 0, cap(carrier.Cache))
-	var packageData []byte
+// 	var packageData []byte
-	var _rest []byte
+// 	var _rest []byte
-	cache, _ := carrier.Cache.Get()
+// 	cache, _ := carrier.Cache.Get()
-	for {
+// 	for {
-		//首先检查这个是不是完整的包，是就返回好了，免得被阻塞
+// 		//首先检查这个是不是完整的包，是就返回好了，免得被阻塞
-		data, rest, err := UnwrapPackage(wrapedPackage)
+// 		data, rest, err := UnwrapPackage(wrapedPackage)
-		packageData = data
+// 		packageData = data
-		_rest = rest
+// 		_rest = rest
-		if err, ok := err.(*NotPackageError); len(wrapedPackage) >= 18 && ok {
+// 		if err, ok := err.(*NotPackageError); len(wrapedPackage) >= 18 && ok {
-			log.Debug("return NotPackageError %s", carrier.AttachedTunnelID)
+// 			log.Debug("return NotPackageError %s", carrier.AttachedTunnelID)
-			return 0, err
+// 			return 0, err
-		}
+// 		}
-		if err == nil {
+// 		if err == nil {
-			//够一个完整的包
+// 			//够一个完整的包
-			capBuff := cap(carrier.GetReceiveBuff())
+// 			capBuff := cap(carrier.GetReceiveBuff())
-			_buff := make([]byte, 0, capBuff) //释放
+// 			_buff := make([]byte, 0, capBuff) //释放
-			_buff = append(_buff, _rest...)
+// 			_buff = append(_buff, _rest...)
-			carrier.SetReceiveBuff(_buff)
+// 			carrier.SetReceiveBuff(_buff)
-			break
+// 			break
-		}
+// 		}
-		//如果读到的数据不够一个完整的包
+// 		//如果读到的数据不够一个完整的包
-		if len(wrapedPackage) > 0 {
+// 		if len(wrapedPackage) > 0 {
-			n, err = carrier.Conn.Read(cache)
+// 			n, err = carrier.Conn.Read(cache)
-			if err != nil {
+// 			if err != nil {
-				log.Error("ERROR %s", err)
+// 				log.Error("ERROR %s", err)
-			}
+// 			}
-		} else {
+// 		} else {
-			n, err = io.ReadAtLeast(carrier.Conn, cache, 18)
+// 			n, err = io.ReadAtLeast(carrier.Conn, cache, 18)
-		}
+// 		}
-
+//
-		if err != nil {
+// 		if err != nil {
-			n = 0
+// 			n = 0
-			return n, err
+// 			return n, err
-		}
+// 		}
-		wrapedPackage = append(wrapedPackage, cache[:n]...)
+// 		wrapedPackage = append(wrapedPackage, cache[:n]...)
-	}
+// 	}
-	decrypted, err := carrier.Cipher.Decrypt(packageData)
+// 	decrypted, err := carrier.Cipher.Decrypt(packageData)
-	if err != nil {
+// 	if err != nil {
-		n = 0
+// 		n = 0
-		return
+// 		return
-	}
+// 	}
-	n = len(decrypted)
+// 	n = len(decrypted)
-	carrier.Msg.Put(decrypted, n)
+// 	carrier.Msg.Put(decrypted, n)
-	return
+// 	return
-}
+// }