update

shadow/packer.go

@@ -3,12 +3,14 @@ package shadow
 import (
 	"bytes"
 	"encoding/binary"
+	"errors"
 	"fmt"
+	"io"
 	"net"
 )
 
 var (
-	AppDataHeader = []byte("GGGGGGGG")
+	AppDataHeader = []byte{0x17, 0x3, 0x3}
 	HeaderLength  = len(AppDataHeader)
 )
 
@@ -17,25 +19,48 @@ type PackAppData struct {
 }
 
 func (m PackAppData) Read(p []byte) (n int, err error) {
-	buf := make([]byte, 32*1024)
-	read, err := m.Conn.Read(buf[0 : HeaderLength+2])
+
+	buf := make([]byte, 32*1024+HeaderLength+2)
+
+	headRead, err := io.ReadAtLeast(m.Conn, buf[0:HeaderLength+2], HeaderLength+2)
+	if err != nil {
+		fmt.Printf("Read header error: %v\n", err)
+		return 0, err
+	}
+	if headRead < HeaderLength+2 {
+		return 0, errors.New("Read header failed")
+	}
 	if bytes.Equal(buf[0:HeaderLength], AppDataHeader) {
-		u := int(binary.BigEndian.Uint16(buf[HeaderLength : HeaderLength+2]))
-		r, err := m.Conn.Read(buf[HeaderLength+2 : u+HeaderLength+2])
-		copy(p, buf[HeaderLength+2:r+HeaderLength+2])
-		return r, err
+		payLoadLength := int(binary.BigEndian.Uint16(buf[HeaderLength : HeaderLength+2]))
+		sum := 0
+		for sum < payLoadLength {
+			r, e := m.Conn.Read(buf[HeaderLength+2+sum : HeaderLength+2+payLoadLength])
+			if e != nil {
+				if e == io.EOF {
+					break
+				} else {
+					return 0, e
+				}
+			}
+			copy(p[sum:], buf[HeaderLength+2+sum:HeaderLength+2+sum+r])
+			sum += r
+		}
+		return sum, err
 	} else {
-		fmt.Println("Header is not present")
-		return read, err
+		fmt.Printf("Invalid header")
+		return 0, errors.New("invalid header")
 	}
 }
 
 func (m PackAppData) Write(p []byte) (n int, err error) {
-	t := make([]byte, len(p)+HeaderLength+2)
-	copy(t[0:], AppDataHeader)
-	binary.BigEndian.PutUint16(t[HeaderLength:], uint16(len(p)))
-	copy(t[HeaderLength+2:], p)
-	write, err := m.Conn.Write(t)
+	lenNum := make([]byte, 2)
+	binary.BigEndian.PutUint16(lenNum, uint16(len(p)))
+
+	packetBuf := bytes.NewBuffer(AppDataHeader)
+	packetBuf.Write(lenNum)
+	packetBuf.Write(p)
+
+	write, err := m.Conn.Write(packetBuf.Bytes())
 	write = write - HeaderLength - 2
 	return write, err
 }