From c2ab6cbe5fbccad47e57091bc0e56174a6050713 Mon Sep 17 00:00:00 2001
From: wenyifan <evan.wen@movit-tech.com>
Date: Thu, 8 Sep 2022 00:12:01 +0800
Subject: [PATCH] update README.md

---
 shadow/client_test.go |  4 ++--
 shadow/packer.go      |  4 +++-
 shadow/server.go      | 30 +++++++++++++-----------------
 shadow/tls_util.go    | 16 +++++++++++-----
 4 files changed, 29 insertions(+), 25 deletions(-)

diff --git a/shadow/client_test.go b/shadow/client_test.go
index d4ac593..45c68cb 100644
--- a/shadow/client_test.go
+++ b/shadow/client_test.go
@@ -11,8 +11,8 @@ import (
 func TestName(t *testing.T) {
 	dial, err := tls.DialWithDialer(&net.Dialer{
 		Timeout: time.Second * 5,
-	}, "tcp", "www.baidu.com:443", &tls.Config{
-		ServerName: "www.baidu.com",
+	}, "tcp", "evan.run:443", &tls.Config{
+		ServerName: "evan.run",
 	})
 
 	err = dial.Handshake()
diff --git a/shadow/packer.go b/shadow/packer.go
index f8b7608..9af56d7 100644
--- a/shadow/packer.go
+++ b/shadow/packer.go
@@ -24,7 +24,9 @@ func (m PackAppData) Read(p []byte) (n int, err error) {
 
 	headRead, err := io.ReadAtLeast(m.Conn, buf[0:HeaderLength+2], HeaderLength+2)
 	if err != nil {
-		fmt.Printf("Read header error: %v\n", err)
+		if err != io.EOF {
+			fmt.Printf("Read header error: %v\n", err)
+		}
 		return 0, err
 	}
 	if headRead < HeaderLength+2 {
diff --git a/shadow/server.go b/shadow/server.go
index f76cb6b..e589575 100644
--- a/shadow/server.go
+++ b/shadow/server.go
@@ -51,14 +51,15 @@ func handler(conn net.Conn, targetAddress string, fakeAddress string) {
 	}
 	waitCh := make(chan int, 1)
 
-	go processHandshake(conn, fakeConn, waitCh)
-	go processHandshake(fakeConn, conn, waitCh)
+	go processHandshake(conn, fakeConn, waitCh, "client")
+	go processHandshake(fakeConn, conn, waitCh, "server")
 
 	<-waitCh
 
 	//Clean up previous buffered data
 	conn.SetDeadline(time.Now())
 	conn.SetDeadline(time.Time{})
+	fakeConn.Close()
 
 	realConnection, err := net.Dial("tcp", targetAddress)
 	if err != nil {
@@ -132,26 +133,22 @@ func handler(conn net.Conn, targetAddress string, fakeAddress string) {
 	//}()
 }
 
-func processHandshake(src net.Conn, dst net.Conn, waitCh chan int) {
+func processHandshake(src net.Conn, dst net.Conn, waitCh chan int, srcType string) {
 	buf := make([]byte, 32*1024)
 	for {
 		nr, er := src.Read(buf)
 		if nr > 0 {
-
-			header := ParseAndVerifyTLSHeader(buf[0:nr])
 			nw, ew := dst.Write(buf[0:nr])
-			if header != nil && header.Type == ChangeCipherSpec {
-				//fmt.Println(header.toString())
-				fmt.Println("[Server] handshake complete")
-				if header.ChangeCipherSpecNext == AppData {
-					dst.Close()
-					waitCh <- 1
-				} else {
-					src.Close()
-					waitCh <- 1
-					return
+			if srcType == "client" {
+				header := ParseAndVerifyTLSHeader(buf[0:nr])
+				if header != nil {
+					if header.Type == ChangeCipherSpec {
+						fmt.Println("[Server] handshake complete")
+						waitCh <- 1
+						break
+					}
+					//fmt.Println(header.toString())
 				}
-				break
 			}
 			if nw < 0 || nr < nw {
 				nw = 0
@@ -175,5 +172,4 @@ func processHandshake(src net.Conn, dst net.Conn, waitCh chan int) {
 			break
 		}
 	}
-	waitCh <- 1
 }
diff --git a/shadow/tls_util.go b/shadow/tls_util.go
index e325372..dd68416 100644
--- a/shadow/tls_util.go
+++ b/shadow/tls_util.go
@@ -17,8 +17,12 @@ const (
 	VersionTLS12 = 0x0303
 	VersionTLS13 = 0x0304
 
-	ServerHello = 2
-	ClientHello = 1
+	ServerHello        = 2
+	ClientHello        = 1
+	Certificate        = 11
+	ServerKeyExchange  = 12
+	ServerHelloDone    = 14
+	EncryptedHandshake = 99
 )
 
 type TLSHeader struct {
@@ -88,12 +92,14 @@ func ParseAndVerifyTLSHeader(data []byte) *TLSHeader {
 	if header.Type == Handshake {
 		header.HandshakeType = data[5]
 		//Check Handshake type
-		if header.HandshakeType != ServerHello && header.HandshakeType != ClientHello {
-			return nil
+		if header.HandshakeType != ServerHello && header.HandshakeType != ClientHello && header.HandshakeType != Certificate && header.HandshakeType != ServerKeyExchange && header.HandshakeType != ServerHelloDone {
+			header.HandshakeType = EncryptedHandshake
 		}
 	}
 	if header.Type == ChangeCipherSpec {
-		header.ChangeCipherSpecNext = data[6]
+		if len(data) > 6 {
+			header.ChangeCipherSpecNext = data[6]
+		}
 	}
 	return header
 }