From fa064736a89b4add380b63581765d7c07be29975 Mon Sep 17 00:00:00 2001
From: LiquidTheDangerous <semenov.andrey.2003@gmail.com>
Date: Wed, 14 Aug 2024 02:33:58 +0500
Subject: [PATCH] refactor

---
 proxy/client_hello.go | 44 ++++++++++++++++++++++++++++++-------------
 proxy/https.go        |  6 +++---
 2 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/proxy/client_hello.go b/proxy/client_hello.go
index 102b7c8..09d1f57 100644
--- a/proxy/client_hello.go
+++ b/proxy/client_hello.go
@@ -7,28 +7,39 @@ import (
 
 const headerLen = 5
 
-type ClientHello struct {
-	Header     ClientHelloHeader
+type TLSMessageType byte
+
+const (
+	TLSInvalid          TLSMessageType = 0x0
+	TLSChangeCipherSpec TLSMessageType = 0x14
+	TLSAlert            TLSMessageType = 0x15
+	TLSHandshake        TLSMessageType = 0x16
+	TLSApplicationData  TLSMessageType = 0x17
+	TLSHeartbeat        TLSMessageType = 0x18
+)
+
+type TlsMessage struct {
+	Header     TlsHeader
 	Raw        []byte //Header + Payload
 	RawHeader  []byte
 	RawPayload []byte
 }
 
-type ClientHelloHeader struct {
-	Type         byte
-	ProtoVersion uint16
+type TlsHeader struct {
+	Type         TLSMessageType
+	ProtoVersion uint16 // major | minor
 	PayloadLen   uint16
 }
 
-func ReadClientHello(r io.Reader) (*ClientHello, error) {
+func ReadTlsMessage(r io.Reader) (*TlsMessage, error) {
 	var rawHeader [5]byte
 	_, err := io.ReadFull(r, rawHeader[:])
 	if err != nil {
 		return nil, err
 	}
 
-	header := ClientHelloHeader{
-		Type:         rawHeader[0],
+	header := TlsHeader{
+		Type:         TLSMessageType(rawHeader[0]),
 		ProtoVersion: binary.BigEndian.Uint16(rawHeader[1:3]),
 		PayloadLen:   binary.BigEndian.Uint16(rawHeader[3:5]),
 	}
@@ -38,11 +49,18 @@ func ReadClientHello(r io.Reader) (*ClientHello, error) {
 	if err != nil {
 		return nil, err
 	}
-	hello := &ClientHello{
-		Header: header,
-		Raw:    raw,
+	hello := &TlsMessage{
+		Header:     header,
+		Raw:        raw,
+		RawHeader:  raw[:headerLen],
+		RawPayload: raw[headerLen:],
 	}
-	hello.RawHeader = hello.Raw[:headerLen]
-	hello.RawPayload = hello.Raw[headerLen:]
 	return hello, nil
 }
+
+func IsClientHello(message *TlsMessage) bool {
+	// According to RFC 8446 section 4.
+	// first byte (Raw[5]) of handshake message should be 0x1 - means client_hello
+	return message.Header.Type == TLSHandshake &&
+		message.Raw[5] == 0x1
+}
diff --git a/proxy/https.go b/proxy/https.go
index 1ef49f8..94eb1c4 100644
--- a/proxy/https.go
+++ b/proxy/https.go
@@ -45,12 +45,12 @@ func (pxy *Proxy) handleHttps(lConn *net.TCPConn, exploit bool, initPkt *packet.
 	log.Debug("[HTTPS] Sent 200 Connection Estabalished to ", lConn.RemoteAddr())
 
 	// Read client hello
-	hello, err := ReadClientHello(lConn)
-	if err != nil {
+	m, err := ReadTlsMessage(lConn)
+	if err != nil || !IsClientHello(m) {
 		log.Debug("[HTTPS] Error reading client hello from the client", err)
 		return
 	}
-	clientHello := hello.Raw
+	clientHello := m.Raw
 
 	log.Debug("[HTTPS] Client sent hello ", len(clientHello), "bytes")