From 26b640e83d48bddb3dd0ad41702e715730cbd159 Mon Sep 17 00:00:00 2001
From: BUPTCZQ <9695911+buptczq@users.noreply.github.com>
Date: Thu, 25 Mar 2021 19:36:52 +0800
Subject: [PATCH] fix: improve the compatibility of XShell

---
 app/xshell.go    | 43 ++++++++++++++++++++++++++++++++++++-------
 versioninfo.json |  6 +++---
 2 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/app/xshell.go b/app/xshell.go
index aa28b32..f4a3de2 100644
--- a/app/xshell.go
+++ b/app/xshell.go
@@ -1,6 +1,7 @@
 package app
 
 import (
+	"bytes"
 	"context"
 	"encoding/binary"
 	"fmt"
@@ -46,7 +47,7 @@ func (s *XShell) Run(ctx context.Context, handler func(conn io.ReadWriteCloser))
 		}
 		wg.Add(1)
 		go func(c io.ReadWriteCloser) {
-			w := &xshellProxy{c, nil}
+			w := &xshellProxy{conn: c}
 			handler(w)
 			wg.Done()
 		}(conn)
@@ -110,19 +111,26 @@ func xshellHandshake(conn net.Conn, cookie string) error {
 	var repMsg initAgentRepMsg
 	repMsg.Flag = initMsg.Flag
 	rep := ssh.Marshal(&repMsg)
-	binary.BigEndian.PutUint32(length[:], uint32(len(rep)))
-	if _, err := conn.Write(length[:]); err != nil {
+	buf := bytes.NewBuffer(nil)
+	err := binary.Write(buf, binary.BigEndian, uint32(len(rep)))
+	if err != nil {
+		return err
+	}
+	_, err = buf.Write(rep)
+	if err != nil {
 		return err
 	}
-	if _, err := conn.Write(rep); err != nil {
+	if _, err := conn.Write(buf.Bytes()); err != nil {
 		return err
 	}
 	return nil
 }
 
 type xshellProxy struct {
-	conn io.ReadWriteCloser
-	buf  []byte
+	conn    io.ReadWriteCloser
+	buf     []byte
+	wlength int
+	wbuf    []byte
 }
 
 type signRequestAgentMsg struct {
@@ -168,7 +176,28 @@ func (s *xshellProxy) Read(p []byte) (n int, err error) {
 }
 
 func (s *xshellProxy) Write(p []byte) (n int, err error) {
-	return s.conn.Write(p)
+	// xshell treats TCP as a message-oriented connection
+	// this piece of sh*t code is in order to be compatible with xshell
+	if s.wlength == 0 {
+		if len(p) != 4 {
+			return 0, fmt.Errorf("xagent proxy: invalid write status")
+		}
+		s.wlength = int(binary.BigEndian.Uint32(p)) + 4
+		s.wbuf = append(s.wbuf, p...)
+	} else {
+		s.wbuf = append(s.wbuf, p...)
+		if len(s.wbuf) == s.wlength {
+			s.wlength = 0
+			_, err := s.conn.Write(s.wbuf)
+			if err != nil {
+				return 0, err
+			}
+			s.wbuf = nil
+		} else if len(s.wbuf) > s.wlength {
+			return 0, fmt.Errorf("xagent proxy: invalid write length")
+		}
+	}
+	return len(p), nil
 }
 
 func (s *xshellProxy) Close() error {
diff --git a/versioninfo.json b/versioninfo.json
index 6dd91aa..1af6c3a 100644
--- a/versioninfo.json
+++ b/versioninfo.json
@@ -3,13 +3,13 @@
         "FileVersion": {
             "Major": 1,
             "Minor": 1,
-            "Patch": 7,
+            "Patch": 8,
             "Build": 0
         },
         "ProductVersion": {
             "Major": 1,
             "Minor": 1,
-            "Patch": 7,
+            "Patch": 8,
             "Build": 0
         },
         "FileFlagsMask": "3f",
@@ -29,7 +29,7 @@
         "OriginalFilename": "WinCryptSSHAgent.exe",
         "PrivateBuild": "",
         "ProductName": "WinCrypt SSH Agent",
-        "ProductVersion": "v1.1.7",
+        "ProductVersion": "v1.1.8",
         "SpecialBuild": ""
     },
     "VarFileInfo": {