消除一个线程引起的bug

Signed-off-by: dmy@lab <dmy@lab.com>
2015-10-18 17:44:56 +08:00 · 2015-10-18 17:44:56 +08:00 · b97df11243
parent 3be197ebba
commit b97df11243
5 changed files with 64 additions and 36 deletions
--- a/cli/cli.go
+++ b/cli/cli.go
@ -7,9 +7,11 @@ import (
 var DestIP string
 var DestPort int
 var ListenPort int
+var EncryptSide string

 func init() {
 	flag.StringVar(&DestIP, "destip", "", "Destination IP")
-	flag.IntVar(&DestPort, "desport", 0, "Destination Port")
+	flag.StringVar(&EncryptSide, "encrypt", "", "Encrypt Side")
+	flag.IntVar(&DestPort, "destport", 0, "Destination Port")
 	flag.IntVar(&ListenPort, "listenport", 0, "Listen Port")
 }
--- a/main.go
+++ b/main.go
@ -11,11 +11,12 @@ import (

 func tunnel() {
 	trans := NewTransTCP()
+	log.Info("%s side is encrypted.", cli.EncryptSide)
 	if cli.DestPort != 0 {
-		log.Info("Listening on 127.0.0.1:%d. Forward %s:%d",cli.ListenPort, cli.DestIP, cli.DestPort)
-		trans.Start(strconv.Itoa(cli.ListenPort), cli.DestIP, strconv.Itoa(cli.DestPort))
+		log.Info("Listening on 127.0.0.1:%d. Forward %s:%d", cli.ListenPort, cli.DestIP, cli.DestPort)
+		trans.Start(strconv.Itoa(cli.ListenPort), cli.DestIP, strconv.Itoa(cli.DestPort), cli.EncryptSide)
 	} else {
-		trans.Start("1200", "192.168.0.120", "8118")
+		trans.Start("1200", "192.168.0.120", "8118", "client")
 	}

 }
--- a/main_test.go
+++ b/main_test.go
@ -2,7 +2,6 @@ package main

 import (
 	"github.com/TransX/log"
-	"github.com/TransX/tscipher"
 	"net"
 	"testing"
 	"time"
@ -20,13 +19,9 @@ func server(t *testing.T) {
 			t.Fatal(err)
 		}
 		bytes := make([]byte, 32)
-		XOR := tscipher.NewXOR([]byte("fasdfasdf"))
 		n, err := conn.Read(bytes)
-		decryped, _ := XOR.Decrypt(bytes[:n])
-		copy(bytes, decryped[:n])
 		log.Info("Test Server Receive %s", string(bytes[:n]))
-		encrypted, _ := XOR.Encrypt([]byte("OK"))
-		_, err = conn.Write(encrypted)
+		_, err = conn.Write([]byte("OK"))
 		log.Info("Test Server write")
 		if err != nil {
 			t.Fatal(err)
@ -42,19 +37,15 @@ func client(t *testing.T) {
 		if err != nil {
 			t.Fatal(err)
 		}
-		XOR := tscipher.NewXOR([]byte("fasdfasdf"))
-		encrypted, _ := XOR.Encrypt([]byte("Client"))
-		conn.Write(encrypted)
-		log.Info("Test Client write %s",string(encrypted))
+		conn.Write([]byte("Client"))
+		log.Info("Test Client write")
 		bytes := make([]byte, 32)
 		n, err := conn.Read(bytes)
-		decryped, _ := XOR.Decrypt(bytes[:n])
-		copy(bytes, decryped[:n])
 		log.Info("Test Client read")
 		if err != nil {
 			t.Fatal(err)
 		}
-		log.Info("Test Client Receive")
+		log.Info("Test Client Receive %s", bytes[:n])
 		time.Sleep(time.Second * 2)
 		conn.Close()
 		log.Info("Test Client closed")
@ -63,11 +54,12 @@ func client(t *testing.T) {
 }

 func TestTunnel(t *testing.T) {
-	log.LogTo("stdout", "INFO")
-	//	t.Log("Start testing.")
+	log.LogTo("log.txt", "DEBUG")
 	log.Info("Test Start testing.")
 	go server(t)
 	go client(t)
-	trans := NewTransTCP()
-	trans.Start("1200", "127.0.0.1", "1244")
+	trans1 := NewTransTCP()
+	go trans1.Start("1200", "127.0.0.1", "1201", "client")
+	trans2 := NewTransTCP()
+	trans2.Start("1201", "127.0.0.1", "1244", "server")
 }
--- a/tcp.go
+++ b/tcp.go
@ -13,15 +13,18 @@ import (
 	"time"
 )

+var seed int32
+
+func init() {
+	seed = 0
+}
+
 type TransTCP struct {
-	seed int32
 }

 func NewTransTCP() *TransTCP {

-	return &TransTCP{
-		0,
-	}
+	return &TransTCP{}
 }

 func (this *TransTCP) createTCPClient(ip, port string) (conn net.Conn, err error) {
@ -47,7 +50,7 @@ func (this *TransTCP) createTCPListener(ip, port string) (listen net.Listener, e
 	return
 }

-func (this *TransTCP) tunnel(src, dest net.Conn, id string) {
+func (this *TransTCP) tunnel(src, dest net.Conn, id string, encrypDirection string) { //单向的，从src发送到dest
 	defer func() {
 		if r := recover(); r != nil {
 			if src != nil {
@ -67,19 +70,29 @@ func (this *TransTCP) tunnel(src, dest net.Conn, id string) {
 			tscipher.NewCipher("XOR"),
 			cache,
 		}
+		log.Debug("Encrypt Direction %s ID %s", encrypDirection, id)
+		if encrypDirection != "receive" {
+			revCarrier.Cipher = nil
+			log.Debug("Read not crypted. Tunnel: %s", id)
+		}
 		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.Info("Reived %d bytes. Tunnel: id %s", nByte,id)
-		log.Debug("Reived %s",cache[:nByte])
+		log.Info("Reived %d bytes. Tunnel: id %s", nByte, id)
+		log.Debug("Reived %s %s", id, cache[:nByte])
 		sendCarrier := &tscipher.Carrier{
 			dest,
 			tscipher.NewCipher("XOR"),
 			cache, //TODO:危险，cache的容量容易被不小心修改
 		}
+		if encrypDirection != "send" {
+			sendCarrier.Cipher = nil
+			log.Debug("Write not crypted. Tunnel: %s", id)
+		}
 		_, err = tscipher.SendData(sendCarrier, nByte)
-		log.Info("Write %d bytes. Tunnel: %s",nByte,id)
+		log.Info("Write %d bytes. Tunnel: %s", nByte, id)
+		log.Debug("Write %s %s", id, cache[:nByte])
 		if err != nil {
 			log.Panic("Write panic. ID: %s, Err: %s, Remote Add: %s", id, err, dest.RemoteAddr().String())
 		}
@ -88,14 +101,14 @@ func (this *TransTCP) tunnel(src, dest net.Conn, id string) {
 }

 func (this *TransTCP) tunnelID() string {
-	nowString := time.Now().String() + strconv.Itoa(int(this.seed))
-	atomic.AddInt32(&this.seed, 1) //避免多线程情况下获得的种子相同
+	nowString := time.Now().String() + strconv.Itoa(int(seed))
+	atomic.AddInt32(&seed, 1) //避免多线程情况下获得的种子相同
 	md5Byte := md5.Sum(bytes.NewBufferString(nowString).Bytes())
+	//	log.Info("seed %d %s", seed, hex.EncodeToString(md5Byte[:]))
 	return hex.EncodeToString(md5Byte[:])
 }

-func (this *TransTCP) Start(listenPort, destIP, destPort string) {
-
+func (this *TransTCP) Start(listenPort, destIP, destPort string, clientOrServer string) {
 	listener, err := this.createTCPListener("0.0.0.0", listenPort)
 	if err != nil {
 		log.Panic("Failed to create listener. %s", err)
@ -111,8 +124,20 @@ func (this *TransTCP) Start(listenPort, destIP, destPort string) {
 				os.Exit(0)
 			}
 			log.Info("Dial %s", destConn.RemoteAddr().String())
-			go this.tunnel(listenerConn, destConn, this.tunnelID())
-			go this.tunnel(destConn, listenerConn, this.tunnelID())
+			//tunnel model : [ -->>server ---- client -->> ](this is a tunnel)
+			tunnelIDA := this.tunnelID()
+			tunnelIDB := this.tunnelID()
+			if clientOrServer == "client" {
+				go this.tunnel(listenerConn, destConn, tunnelIDA, "send")
+				go this.tunnel(destConn, listenerConn, tunnelIDB, "receive")
+				log.Debug("tow tunnel created: %s %s %s %s", tunnelIDA, "send", tunnelIDB, "receive")
+			}
+			if clientOrServer == "server" {
+				go this.tunnel(listenerConn, destConn, tunnelIDA, "receive")
+				go this.tunnel(destConn, listenerConn, tunnelIDB, "send")
+				log.Debug("tow tunnel created: %s %s %s %s", tunnelIDA, "receive", tunnelIDB, "send")
+			}
+
 		} else {
 			log.Info("Failed to accept incoming connection. %s", err)
 		}
--- a/tscipher/cipher.go
+++ b/tscipher/cipher.go
@ -23,18 +23,23 @@ func NewCipher(cipherName string) (cipher Cipher) {
 		return NewAES()
 	}
 	if cipherName == "XOR" {
-		return NewXOR([]byte("fasdfasdf"))
+		return NewXOR([]byte("fasdfasdf!3297!jfsl12*&!HHHFds"))
 	}
 	return nil //TODO:临时这样处理
 }

 func SendData(carrier *Carrier, nByte int) (n int, err error) {
+	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
 	}
 	n, err = carrier.Conn.Write(encrypedByte[:nByte])
+	copy(carrier.Cache, encrypedByte[:nByte]) // in case of debugging
 	return
 }

@ -44,6 +49,9 @@ func ReceiveData(carrier *Carrier) (n int, err error) {
 		n = 0
 		return
 	}
+	if carrier.Cipher == nil {
+		return
+	}
 	decrypted, err := carrier.Cipher.Decrypt(carrier.Cache[:n])
 	copy(carrier.Cache, decrypted[:n])
 	if err != nil {