From 8165adcab15347bd1c6613d9ac9cd6bb8b25bdcb Mon Sep 17 00:00:00 2001
From: Alex X <alexey.khit@gmail.com>
Date: Thu, 12 Oct 2023 17:03:58 +0300
Subject: [PATCH] Rewrite hap secure connection

---
 pkg/hap/client.go        |   2 +-
 pkg/hap/secure/secure.go | 109 +++++++++++++++++++++------------------
 2 files changed, 61 insertions(+), 50 deletions(-)

diff --git a/pkg/hap/client.go b/pkg/hap/client.go
index de038b67..9fb5ef3a 100644
--- a/pkg/hap/client.go
+++ b/pkg/hap/client.go
@@ -217,7 +217,7 @@ func (c *Client) Dial() (err error) {
 		return
 	}
 	// new reader for new conn
-	c.reader = bufio.NewReaderSize(c.Conn, 32*1024) // 32K like default request body
+	c.reader = bufio.NewReader(c.Conn)
 
 	go c.eventsReader()
 
diff --git a/pkg/hap/secure/secure.go b/pkg/hap/secure/secure.go
index 29c17ce6..d3d60e5e 100644
--- a/pkg/hap/secure/secure.go
+++ b/pkg/hap/secure/secure.go
@@ -1,6 +1,7 @@
 package secure
 
 import (
+	"bufio"
 	"encoding/binary"
 	"io"
 	"net"
@@ -14,6 +15,10 @@ import (
 type Conn struct {
 	conn net.Conn
 
+	rd *bufio.Reader
+	wr *bufio.Writer
+	rb []byte // temporary reading buffer
+
 	encryptKey []byte
 	decryptKey []byte
 	encryptCnt uint64
@@ -33,11 +38,19 @@ func Client(conn net.Conn, sharedKey []byte, isClient bool) (net.Conn, error) {
 		return nil, err
 	}
 
+	c := &Conn{
+		conn: conn,
+		rd:   bufio.NewReaderSize(conn, BufferSize),
+		wr:   bufio.NewWriterSize(conn, BufferSize),
+	}
+
 	if isClient {
-		return &Conn{conn: conn, encryptKey: key2, decryptKey: key1}, nil
+		c.encryptKey, c.decryptKey = key2, key1
 	} else {
-		return &Conn{conn: conn, encryptKey: key1, decryptKey: key2}, nil
+		c.encryptKey, c.decryptKey = key1, key2
 	}
+
+	return c, nil
 }
 
 const (
@@ -47,87 +60,85 @@ const (
 	VerifySize = 2
 	NonceSize  = 8
 	Overhead   = 16 // chacha20poly1305.Overhead
+
+	BufferSize = 2 + 0xFFFF + Overhead
 )
 
 func (c *Conn) Read(b []byte) (n int, err error) {
-	verify := make([]byte, VerifySize) // = packet length
-	buf := make([]byte, PacketSizeMax+Overhead)
-	nonce := make([]byte, NonceSize)
-
-	for {
-		if len(b) < PacketSizeMax {
-			return
-		}
+	// something in reading buffer
+	if len(c.rb) > 0 {
+		n = copy(b, c.rb)
+		c.rb = c.rb[n:]
+		return
+	}
 
-		if _, err = io.ReadFull(c.conn, verify); err != nil {
-			return
-		}
+	verify := make([]byte, 2) // verify = plain message size
+	if _, err = io.ReadFull(c.rd, verify); err != nil {
+		return
+	}
 
-		size := binary.LittleEndian.Uint16(verify)
-		ciphertext := buf[:size+Overhead]
+	n = int(binary.LittleEndian.Uint16(verify))
+	ciphertext := make([]byte, n+Overhead)
 
-		if _, err = io.ReadFull(c.conn, ciphertext); err != nil {
-			return
-		}
+	if _, err = io.ReadFull(c.rd, ciphertext); err != nil {
+		return
+	}
 
-		binary.LittleEndian.PutUint64(nonce, c.decryptCnt)
-		c.decryptCnt++
+	nonce := make([]byte, NonceSize)
+	binary.LittleEndian.PutUint64(nonce, c.decryptCnt)
+	c.decryptCnt++
 
-		// put decrypted text to b's end
+	// buffer size enought for direct reading
+	if n <= cap(b) {
 		_, err = chacha20poly1305.DecryptAndVerify(c.decryptKey, b[:0], nonce, ciphertext, verify)
-		if err != nil {
-			return
-		}
-
-		n += int(size) // plaintext size
-
-		// Finish when all bytes fit in b
-		if size < PacketSizeMax {
-			return
-		}
+		return
+	}
 
-		b = b[size:]
+	// reading to temporary buffer
+	c.rb = make([]byte, 0, n)
+	if _, err = chacha20poly1305.DecryptAndVerify(c.decryptKey, c.rb, nonce, ciphertext, verify); err != nil {
+		return
 	}
+	return c.Read(b)
 }
 
 func (c *Conn) Write(b []byte) (n int, err error) {
-	c.mx.Lock()
-	defer c.mx.Unlock()
-
-	nonce := make([]byte, NonceSize)
-	buf := make([]byte, NonceSize+PacketSizeMax+Overhead)
-	verify := buf[:VerifySize] // part of write buffer
+	var ciphertext []byte
+	var nonce = make([]byte, NonceSize)
+	var verify = make([]byte, VerifySize)
 
-	for {
+	for len(b) > 0 {
 		size := len(b)
 		if size > PacketSizeMax {
 			size = PacketSizeMax
 		}
 
 		binary.LittleEndian.PutUint16(verify, uint16(size))
+		if _, err = c.wr.Write(verify); err != nil {
+			return
+		}
 
 		binary.LittleEndian.PutUint64(nonce, c.encryptCnt)
 		c.encryptCnt++
 
-		// put encrypted text to writing buffer just after size (2 bytes)
-		_, err = chacha20poly1305.EncryptAndSeal(c.encryptKey, buf[2:2], nonce, b[:size], verify)
-		if err != nil {
-			return
+		if cap(ciphertext) < size+Overhead {
+			ciphertext = make([]byte, size+Overhead)
 		}
 
-		if _, err = c.conn.Write(buf[:VerifySize+size+Overhead]); err != nil {
+		_, err = chacha20poly1305.EncryptAndSeal(c.encryptKey, ciphertext[:0], nonce, b[:size], verify)
+		if err != nil {
 			return
 		}
 
-		n += size // plaintext size
-
-		if size < PacketSizeMax {
-			break
+		if _, err = c.wr.Write(ciphertext); err != nil {
+			return
 		}
 
-		b = b[PacketSizeMax:]
+		b = b[size:]
+		n += size
 	}
 
+	err = c.wr.Flush()
 	return
 }