X-Git-Url: https://git.r.bdr.sh/rbdr/mobius/blobdiff_plain/aebc4d3647b9823ae8cbb57b21b1af83bfd011fb..270abe9a991a57084650424a0238ef4f062b0b64:/hotline/handshake.go?ds=inline

diff --git a/hotline/handshake.go b/hotline/handshake.go
index 66bf643..9ccdc62 100644
--- a/hotline/handshake.go
+++ b/hotline/handshake.go
@@ -4,7 +4,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"errors"
-	"net"
+	"io"
 )
 
 type handshake struct {
@@ -14,6 +14,8 @@ type handshake struct {
 	SubVersion  [2]byte
 }
 
+var trtp = [4]byte{0x54, 0x52, 0x54, 0x50}
+
 // Handshake
 // 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.
@@ -33,17 +35,22 @@ type handshake struct {
 // Description		Size 	Data	Note
 // Protocol ID		4		TRTP
 // Error code		4				Error code returned by the server (0 = no error)
-func Handshake(conn net.Conn, buf []byte) error {
+func Handshake(rw io.ReadWriter) error {
+	handshakeBuf := make([]byte, 12)
+	if _, err := io.ReadFull(rw, handshakeBuf); err != nil {
+		return err
+	}
+
 	var h handshake
-	r := bytes.NewReader(buf)
+	r := bytes.NewReader(handshakeBuf)
 	if err := binary.Read(r, binary.BigEndian, &h); err != nil {
 		return err
 	}
 
-	if h.Protocol != [4]byte{0x54, 0x52, 0x54, 0x50} {
+	if h.Protocol != trtp {
 		return errors.New("invalid handshake")
 	}
 
-	_, err := conn.Write([]byte{84, 82, 84, 80, 0, 0, 0, 0})
+	_, err := rw.Write([]byte{84, 82, 84, 80, 0, 0, 0, 0})
 	return err
 }