Correctly manage partial writes

This commit is contained in:
9seconds
2021-03-29 10:26:52 +03:00
parent db8614999a
commit 3992054560
7 changed files with 75 additions and 38 deletions
+12 -7
View File
@@ -10,11 +10,9 @@ type Conn struct {
Encryptor cipher.Stream
Decryptor cipher.Stream
writeBuf []byte
}
func (c *Conn) Read(p []byte) (int, error) {
func (c Conn) Read(p []byte) (int, error) {
n, err := c.Conn.Read(p)
if err != nil {
return n, err // nolint: wrapcheck
@@ -25,9 +23,16 @@ func (c *Conn) Read(p []byte) (int, error) {
return n, nil
}
func (c *Conn) Write(p []byte) (int, error) {
c.writeBuf = append(c.writeBuf[:0], p...)
c.Encryptor.XORKeyStream(c.writeBuf, p)
func (c Conn) Write(p []byte) (int, error) {
buf := acquireBytesBuffer()
defer releaseBytesBuffer(buf)
return c.Conn.Write(c.writeBuf)
buf.Write(p)
payload := buf.Bytes()
c.Encryptor.XORKeyStream(payload, payload)
n, err := buf.WriteTo(c.Conn)
return int(n), err // nolint: wrapcheck
}