|
1 | 1 | package app |
2 | 2 |
|
3 | 3 | import ( |
| 4 | + "bytes" |
4 | 5 | "context" |
5 | 6 | "encoding/binary" |
6 | 7 | "fmt" |
@@ -46,7 +47,7 @@ func (s *XShell) Run(ctx context.Context, handler func(conn io.ReadWriteCloser)) |
46 | 47 | } |
47 | 48 | wg.Add(1) |
48 | 49 | go func(c io.ReadWriteCloser) { |
49 | | - w := &xshellProxy{c, nil} |
| 50 | + w := &xshellProxy{conn: c} |
50 | 51 | handler(w) |
51 | 52 | wg.Done() |
52 | 53 | }(conn) |
@@ -110,19 +111,26 @@ func xshellHandshake(conn net.Conn, cookie string) error { |
110 | 111 | var repMsg initAgentRepMsg |
111 | 112 | repMsg.Flag = initMsg.Flag |
112 | 113 | rep := ssh.Marshal(&repMsg) |
113 | | - binary.BigEndian.PutUint32(length[:], uint32(len(rep))) |
114 | | - if _, err := conn.Write(length[:]); err != nil { |
| 114 | + buf := bytes.NewBuffer(nil) |
| 115 | + err := binary.Write(buf, binary.BigEndian, uint32(len(rep))) |
| 116 | + if err != nil { |
| 117 | + return err |
| 118 | + } |
| 119 | + _, err = buf.Write(rep) |
| 120 | + if err != nil { |
115 | 121 | return err |
116 | 122 | } |
117 | | - if _, err := conn.Write(rep); err != nil { |
| 123 | + if _, err := conn.Write(buf.Bytes()); err != nil { |
118 | 124 | return err |
119 | 125 | } |
120 | 126 | return nil |
121 | 127 | } |
122 | 128 |
|
123 | 129 | type xshellProxy struct { |
124 | | - conn io.ReadWriteCloser |
125 | | - buf []byte |
| 130 | + conn io.ReadWriteCloser |
| 131 | + buf []byte |
| 132 | + wlength int |
| 133 | + wbuf []byte |
126 | 134 | } |
127 | 135 |
|
128 | 136 | type signRequestAgentMsg struct { |
@@ -168,7 +176,28 @@ func (s *xshellProxy) Read(p []byte) (n int, err error) { |
168 | 176 | } |
169 | 177 |
|
170 | 178 | func (s *xshellProxy) Write(p []byte) (n int, err error) { |
171 | | - return s.conn.Write(p) |
| 179 | + // xshell treats TCP as a message-oriented connection |
| 180 | + // this piece of sh*t code is in order to be compatible with xshell |
| 181 | + if s.wlength == 0 { |
| 182 | + if len(p) != 4 { |
| 183 | + return 0, fmt.Errorf("xagent proxy: invalid write status") |
| 184 | + } |
| 185 | + s.wlength = int(binary.BigEndian.Uint32(p)) + 4 |
| 186 | + s.wbuf = append(s.wbuf, p...) |
| 187 | + } else { |
| 188 | + s.wbuf = append(s.wbuf, p...) |
| 189 | + if len(s.wbuf) == s.wlength { |
| 190 | + s.wlength = 0 |
| 191 | + _, err := s.conn.Write(s.wbuf) |
| 192 | + if err != nil { |
| 193 | + return 0, err |
| 194 | + } |
| 195 | + s.wbuf = nil |
| 196 | + } else if len(s.wbuf) > s.wlength { |
| 197 | + return 0, fmt.Errorf("xagent proxy: invalid write length") |
| 198 | + } |
| 199 | + } |
| 200 | + return len(p), nil |
172 | 201 | } |
173 | 202 |
|
174 | 203 | func (s *xshellProxy) Close() error { |
|
0 commit comments