Refactoring, cleanup, test backfilling

1 files changed, 52 insertions, 21 deletions

diff --git a/hotline/handshake.go b/hotline/handshake.go
index 9ccdc62..c54359e 100644
--- a/hotline/handshake.go
+++ b/hotline/handshake.go
@@ -4,19 +4,12 @@ import (
 	"bytes"
 	"encoding/binary"
 	"errors"
+	"fmt"
 	"io"
 )
 
-type handshake struct {
-	Protocol    [4]byte // Must be 0x54525450 TRTP
-	SubProtocol [4]byte
-	Version     [2]byte // Always 1
-	SubVersion  [2]byte
-}
-
-var trtp = [4]byte{0x54, 0x52, 0x54, 0x50}
-
-// Handshake
+// Hotline handshake process
+//
 // After establishing TCP connection, both client and server start the handshake process
 // in order to confirm that each of them comply with requirements of the other.
 // The information provided in this initial data exchange identifies protocols,
@@ -35,22 +28,60 @@ var trtp = [4]byte{0x54, 0x52, 0x54, 0x50}
 // Description		Size 	Data	Note
 // Protocol ID		4		TRTP
 // Error code		4				Error code returned by the server (0 = no error)
-func Handshake(rw io.ReadWriter) error {
-	handshakeBuf := make([]byte, 12)
-	if _, err := io.ReadFull(rw, handshakeBuf); err != nil {
-		return err
+
+type handshake struct {
+	Protocol    [4]byte // Must be 0x54525450 TRTP
+	SubProtocol [4]byte // Must be 0x484F544C HOTL
+	Version     [2]byte // Always 1 (?)
+	SubVersion  [2]byte // Always 2 (?)
+}
+
+// Write implements the io.Writer interface for handshake.
+func (h *handshake) Write(p []byte) (n int, err error) {
+	if len(p) != handshakeSize {
+		return 0, errors.New("invalid handshake size")
 	}
 
+	_ = binary.Read(bytes.NewBuffer(p), binary.BigEndian, h)
+
+	return len(p), nil
+}
+
+// Valid checks if the handshake contains valid protocol and sub-protocol IDs.
+func (h *handshake) Valid() bool {
+	return h.Protocol == trtp && h.SubProtocol == hotl
+}
+
+var (
+	// trtp represents the Protocol ID "TRTP" in hex
+	trtp = [4]byte{0x54, 0x52, 0x54, 0x50}
+
+	// hotl represents the Sub-protocol ID "HOTL" in hex
+	hotl = [4]byte{0x48, 0x4F, 0x54, 0x4C}
+
+	// handshakeResponse represents the server's response after a successful handshake
+	// Response with "TRTP" and no error code
+	handshakeResponse = [8]byte{0x54, 0x52, 0x54, 0x50, 0x00, 0x00, 0x00, 0x00}
+)
+
+const handshakeSize = 12
+
+// performHandshake performs the handshake process.
+func performHandshake(rw io.ReadWriter) error {
 	var h handshake
-	r := bytes.NewReader(handshakeBuf)
-	if err := binary.Read(r, binary.BigEndian, &h); err != nil {
-		return err
+
+	// Copy exactly handshakeSize bytes from rw to handshake
+	if _, err := io.CopyN(&h, rw, handshakeSize); err != nil {
+		return fmt.Errorf("failed to read handshake data: %w", err)
+	}
+
+	if !h.Valid() {
+		return errors.New("invalid protocol or sub-protocol in handshake")
 	}
 
-	if h.Protocol != trtp {
-		return errors.New("invalid handshake")
+	if _, err := rw.Write(handshakeResponse[:]); err != nil {
+		return fmt.Errorf("error sending handshake response: %w", err)
 	}
 
-	_, err := rw.Write([]byte{84, 82, 84, 80, 0, 0, 0, 0})
-	return err
+	return nil
 }