facat
/
transx
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

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

This commit is contained in:
dmy@lab 2016-01-31 12:32:41 +08:00
parent dee06848fb
commit bdfcbbbda8
7 changed files with 288 additions and 108 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
benckmark.go
View File

@ -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

104
communicator/receiver.go Normal file
View File

@ -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()
}

79
communicator/sender.go Normal file
View File

@ -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()
}

2
main.go
View File

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

30
model/tunnel.go
View File

@ -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 {
revCarrier.Cipher = nil
n, err = tscipher.RowReceiveData(revCarrier)
} else {
n, err = tscipher.ReceiveData(revCarrier)
}
n, err = receiver.Receive()
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())

12
protocol/tcp.go
View File

@ -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())

167
tscipher/cipher.go
View File

@ -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) {
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 := WrapPackage(encrypedByte[:nByte])
n, err = carrier.Conn.Write(wraped)
carrier.Cache.Put(make([]byte, 1024*4), 1024*4)
return
}
func RowReceiveData(carrier *Carrier) (n int, err error) {
cache, _ := carrier.Cache.Get()
n, err = carrier.Conn.Read(cache)
if err != nil {
n = 0
}
carrier.Msg.Put(cache, n)
return
}
func ReceiveData(carrier *Carrier) (n int, err error) {
wrapedPackage := carrier.GetReceiveBuff() //make([]byte, 0, cap(carrier.Cache))
var packageData []byte
var _rest []byte
cache, _ := carrier.Cache.Get()
for {
//首先检查这个是不是完整的包，是就返回好了，免得被阻塞
data, rest, err := UnwrapPackage(wrapedPackage)
packageData = data
_rest = rest
if err, ok := err.(*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)
}
} 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 SendData(carrier *Carrier) (n int, err error) {
// 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 := WrapPackage(encrypedByte[:nByte])
// n, err = carrier.Conn.Write(wraped)
// carrier.Cache.Put(make([]byte, 1024*4), 1024*4)
// return
// }
//
// func RowReceiveData(carrier *Carrier) (n int, err error) {
// cache, _ := carrier.Cache.Get()
// n, err = carrier.Conn.Read(cache)
// if err != nil {
// n = 0
// }
// carrier.Msg.Put(cache, n)
// return
// }
//
// func ReceiveData(carrier *Carrier) (n int, err error) {
// wrapedPackage := carrier.GetReceiveBuff() //make([]byte, 0, cap(carrier.Cache))
// var packageData []byte
// var _rest []byte
// cache, _ := carrier.Cache.Get()
// for {
// //首先检查这个是不是完整的包，是就返回好了，免得被阻塞
// data, rest, err := UnwrapPackage(wrapedPackage)
// packageData = data
// _rest = rest
// if err, ok := err.(*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)
// }
// } 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
// }