transx/tscipher/cipher.go

package tscipher

import (
	"bytes"
	"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)
}

type Carrier struct {
	Conn             net.Conn
	Cipher           Cipher
	Cache            []byte
	AttachedTunnelID string
	ReceiveBuff      []byte
}

func NewCipher(cipherName string) (cipher Cipher) {
	if cipherName == "default" {
		return NewChaCha()
	}
	if cipherName == "AES" {
		return NewAES()
	}
	if cipherName == "XOR" {
		return NewXOR([]byte("fasdfasdf!3297!jfsl12*&!HHHFds"))
	}
	return nil //TODO:临时这样处理
}

func WrapPackage(data []byte) []byte { //把要加密传输的数据打包成一定的格式，避免发送了100自己，只收到90字节的问题。
	sizeOfData := len(data)
	binSize := utils.Int2binary(sizeOfData, 10)
	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)]
	}
	return append(cipheredHeader, data...)
}

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)]
	}
	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)
		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 = nil
		rest = nil
		err = &NotPackageError{"NotPackageError" + "start:" + string(start) + " end:" + string(end)}
	}
	return

}

func SendData(carrier *Carrier, nByte int) (n int, err error) {
	if len(carrier.Cache) < nByte {
		log.Panic("Cache of send is too small")
	}
	if carrier.Cipher == nil {
		n, err = carrier.Conn.Write(carrier.Cache[:nByte])
		return
	}
	encrypedByte, err := carrier.Cipher.Encrypt(carrier.Cache[:nByte])
	if err != nil {
		n = 0
		return
	}
	//打包
	wraped := WrapPackage(encrypedByte[:nByte])
	n, err = carrier.Conn.Write(wraped)
	copy(carrier.Cache, encrypedByte[:nByte]) // in case of debugging
	return
}

func SendData2(carrier *Carrier, nByte int) (n int, err error) {
	n, err = carrier.Conn.Write(carrier.Cache[:nByte])
	if err != nil {
		return
	}
	return
}

func RowReceiveData(carrier *Carrier) (n int, err error) {
	n, err = carrier.Conn.Read(carrier.Cache)
	if err != nil {
		n = 0

	}
	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)
	copy(carrier.Cache, decrypted)
	return
}

func ReceiveData2(carrier *Carrier) (n int, err error) {
	n, err = carrier.Conn.Read(carrier.Cache)
	if err != nil {
		n = 0
		return
	}
	return
}