diff --git a/main_test.go b/main_test.go index 753d141..6988f3c 100644 --- a/main_test.go +++ b/main_test.go @@ -4,6 +4,7 @@ import ( "bufio" "fmt" "github.com/TransX/log" + "io" "net" "os" "testing" @@ -33,9 +34,9 @@ func serverBin(t *testing.T) { n := 0 for { /////-> - n, err = conn.Read(_bytes) + n, err = io.ReadAtLeast(conn, _bytes, 4097) if err != nil { - log.Error("Test Server read %", err.Error()) + log.Error("Test Server read err %s", err.Error()) break } log.Info("Test Server read per time %d", n) diff --git a/tscipher/cipher.go b/tscipher/cipher.go index 7759cb7..2d385b9 100644 --- a/tscipher/cipher.go +++ b/tscipher/cipher.go @@ -2,15 +2,32 @@ package tscipher import ( "bytes" - "errors" "github.com/TransX/log" "github.com/TransX/utils" + "io" "net" + // "time" ) var StartMark = []byte("#2v!") //should be constant var EndMark = []byte("_=1z") //should be constant +type LackDataError struct { + e string +} + +type NotPackageError struct { + e string +} + +func (this *LackDataError) Error() string { + return this.e +} + +func (this *NotPackageError) Error() string { + return this.e +} + type Cipher interface { Decrypt(data []byte) (decrypted []byte, err error) Encrypt(data []byte) (encryped []byte, err error) @@ -21,6 +38,7 @@ type Carrier struct { Cipher Cipher Cache []byte AttachedTunnelID string + ReceiveBuff []byte } func NewCipher(cipherName string) (cipher Cipher) { @@ -42,32 +60,44 @@ func WrapPackage(data []byte) []byte { //把要加密传输的数据打包成一 header := append(append(StartMark, binSize...), EndMark...) //加密 key := []byte("hahahehe~-1!") + key = key cipheredHeader := make([]byte, len(header)) for i, v := range header { - cipheredHeader[i] = v ^ key[i%len(key)] + cipheredHeader[i] = v //^ key[i%len(key)] } return append(cipheredHeader, data...) } -func UnwrapPackage(pacakge []byte) (packageSize int, data []byte, err error) { +func UnwrapPackage(pacakge []byte) (data []byte, rest []byte, err error) { //前14个字节是header cipheredHeader := pacakge[:18] header := make([]byte, len(cipheredHeader)) key := []byte("hahahehe~-1!") + key = key for i, v := range cipheredHeader { - header[i] = v ^ key[i%len(key)] + header[i] = v //^ key[i%len(key)] } start := header[:4] end := header[14:] binSize := header[4:14] + packageSize := 0 if bytes.Compare(start, StartMark) == 0 && bytes.Compare(end, EndMark) == 0 { packageSize = utils.Binary2Int(binSize) - data = pacakge[18:] + if len(pacakge[18:]) < packageSize { + packageSize = 0 + data = nil + rest = nil + err = &LackDataError{"LackDataError"} + return + } + data = pacakge[18 : 18+packageSize] + rest = pacakge[18+len(data):] err = nil } else { packageSize = 0 - data = pacakge - err = errors.New("not a package") + data = nil + rest = nil + err = &NotPackageError{"NotPackageError" + "start:" + string(start) + " end:" + string(end)} } return @@ -101,56 +131,123 @@ func SendData2(carrier *Carrier, nByte int) (n int, err error) { return } -func ReceiveData(carrier *Carrier) (n int, err error) { +func RowReceiveData(carrier *Carrier) (n int, err error) { n, err = carrier.Conn.Read(carrier.Cache) if err != nil { n = 0 - return - } - if carrier.Cipher == nil { - return - } - //解包 - wrapedPackage := carrier.Cache[:n] - packageSize, data, err := UnwrapPackage(wrapedPackage) - realData := make([]byte, 0, packageSize) - // log.Info("packageSize %d data size %d", packageSize, len(data)) - if err == nil && packageSize == len(data) { //读到的是一个完整的包 - realData = data - n = len(realData) - log.Debug("read a complete package") - } else { - gotSize := len(data) - for { - n, err = carrier.Conn.Read(carrier.Cache) - if err != nil { - n = 0 - return - } - wrapedPackage = carrier.Cache[:n] - log.Debug("got partial package size %d from %s ID: %s", n, carrier.Conn.RemoteAddr().String(), carrier.AttachedTunnelID) - _, data, err = UnwrapPackage(wrapedPackage) - if err == nil { - n = 0 - err = errors.New("partial package lost") - return - } - gotSize += len(data) - realData = append(realData, data...) - if gotSize == packageSize { - log.Debug("got enough:. packageSize %d, real size %d. not include header", packageSize, gotSize) - n = gotSize - break - } - } } - decrypted, err := carrier.Cipher.Decrypt(realData) + return +} + +func ReceiveData(carrier *Carrier) (n int, err error) { + defer func() { + if r := recover(); r != nil { + log.Error("ReceiveData err %s", r) + } + }() + // hasHeader := false + wrapedPackage := carrier.ReceiveBuff //make([]byte, 0, cap(carrier.Cache)) + var packageData []byte + var _rest []byte + for { + //首先检查这个是不是完整的包,是就返回好了,免得被阻塞 + if data, _, err := UnwrapPackage(wrapedPackage); err == nil { + packageData = data + } else { + log.Info("id %s to read wrapedPackage %d", carrier.AttachedTunnelID, len(wrapedPackage)) + // carrier.Conn.SetReadDeadline(time.Now().Add(time.Second * 5)) + if len(wrapedPackage) > 0 { + + n, err = carrier.Conn.Read(carrier.Cache) + if err != nil { + log.Error("ERROR %s", err) + } + log.Info("id %s to Conn.Read %d", carrier.AttachedTunnelID, n) + } else { + n, err = io.ReadAtLeast(carrier.Conn, carrier.Cache, 18) + log.Info("id %s to ReadAtLeast", carrier.AttachedTunnelID) + } + if err != nil { + n = 0 + return n, err + } + + // if len(carrier.ReceiveBuff) > 0 { + // hasHeader = true + // } + //解包 + // if !hasHeader { + // copy(wrapedPackage, carrier.Cache[:n]) + // } else { + // wrapedPackage = append(carrier.ReceiveBuff, carrier.Cache[:n]) + // } + log.Info("length of wrapedPackage %d", len(wrapedPackage)) + wrapedPackage = append(wrapedPackage, carrier.Cache[:n]...) + log.Info("first 18 %s from %s", string(wrapedPackage[:18]), carrier.Conn.RemoteAddr().String()) + data, rest, err := UnwrapPackage(wrapedPackage) + if err, ok := err.(*NotPackageError); ok { + return 0, err + } + //如果读到的数据不够一个完整的包 + if _, ok := err.(*LackDataError); ok { + // carrier.ReceiveBuff = append(carrier.ReceiveBuff, wrapedPackage) + // hasHeader = true + log.Info("not enough") + continue + } + packageData = data + _rest = rest + } + //够一个完整的包 + carrier.ReceiveBuff = make([]byte, 0, cap(carrier.ReceiveBuff)) //释放 + carrier.ReceiveBuff = append(carrier.ReceiveBuff, _rest...) + + break + } + + // realData := make([]byte, 0, packageSize) + // // log.Info("packageSize %d data size %d", packageSize, len(data)) + // if err == nil && packageSize == len(data) { //读到的是一个完整的包 + // realData = data + // n = len(realData) + // log.Debug("read a complete package") + // } else { + // gotSize := len(data) + // for { + // n, err = carrier.Conn.Read(carrier.Cache) + // if err != nil { + // n = 0 + // return + // } + // wrapedPackage = carrier.Cache[:n] + // log.Debug("got partial package size %d from %s ID: %s", n, carrier.Conn.RemoteAddr().String(), carrier.AttachedTunnelID) + // _, data, err = UnwrapPackage(wrapedPackage) + // if err == nil { + // n = 0 + // err = errors.New("partial package lost") + // return + // } + // gotSize += len(data) + // realData = append(realData, data...) + // if gotSize == packageSize { + // log.Debug("got enough:. packageSize %d, real size %d. not include header", packageSize, gotSize) + // n = gotSize + // break + // } + // } + + // } + if len(carrier.ReceiveBuff) > 0 { + log.Info("trailing %d from %s", len(carrier.ReceiveBuff), carrier.Conn.RemoteAddr().String()) + log.Info("18 byte of trailing %s", string(carrier.ReceiveBuff[:18])) + } + decrypted, err := carrier.Cipher.Decrypt(packageData) if err != nil { n = 0 return } - // n = len(decrypted) + n = len(decrypted) copy(carrier.Cache, decrypted) return } diff --git a/tunnel.go b/tunnel.go index 95cc0c2..cd2c4e1 100644 --- a/tunnel.go +++ b/tunnel.go @@ -60,30 +60,37 @@ func (this *Tunnel) run() { //单向的,从src发送到dest } }() cache := make([]byte, 1024*2) //128kB + //构建Carrier + revCarrier := &tscipher.Carrier{ + src, + tscipher.NewCipher("XOR"), + cache, + this.id, + make([]byte, 0, cap(cache)*8), + } + sendCarrier := &tscipher.Carrier{ + dest, + tscipher.NewCipher("XOR"), + cache, //TODO:危险,cache的容量容易被不小心修改 + this.id, + make([]byte, 0, cap(cache)*8), + } for { - //构建Carrier - revCarrier := &tscipher.Carrier{ - src, - tscipher.NewCipher("XOR"), - cache, - this.id, - } + // log.Info("a loop Tunnel id %s", id) + var nByte int // log.Debug("Encrypt Direction %s ID %s", cipherDirection, id) + var err error if cipherDirection != "receive" { revCarrier.Cipher = nil // log.Debug("Read not crypted. Tunnel: %s", id) + nByte, err = tscipher.RowReceiveData(revCarrier) + } else { + nByte, err = tscipher.ReceiveData(revCarrier) } - nByte, err := tscipher.ReceiveData(revCarrier) if err != nil { - log.Panic("Read panic. Tunnel id: %s. Remote Add: %s. Err:%s", id, src.RemoteAddr().String(), err) + log.Panic("Read panic. Tunnel id: %s. Remote Add: %s. Err:%s", id, src.RemoteAddr().String(), err.Error()) } log.Info("Reived %d bytes from %s. Tunnel: id %s", nByte, src.RemoteAddr().String(), id) - sendCarrier := &tscipher.Carrier{ - dest, - tscipher.NewCipher("XOR"), - cache, //TODO:危险,cache的容量容易被不小心修改 - this.id, - } if cipherDirection != "send" { sendCarrier.Cipher = nil // log.Debug("Write not crypted. Tunnel: %s", id)