From 26070d5b3e536e4a4feb10ad8042f81873a3b9e2 Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Thu, 1 Apr 2021 14:12:29 +0300
Subject: [PATCH] More correct calculation of dc for obfuscated2 frame

---
 .../internal/obfuscated2/handshake_frame.go   |  7 +----
 .../handshake_frame_internal_test.go          | 27 +++++++++++++++++++
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/mtglib/internal/obfuscated2/handshake_frame.go b/mtglib/internal/obfuscated2/handshake_frame.go
index fde59ae..fff548f 100644
--- a/mtglib/internal/obfuscated2/handshake_frame.go
+++ b/mtglib/internal/obfuscated2/handshake_frame.go
@@ -1,7 +1,5 @@
 package obfuscated2
 
-import "encoding/binary"
-
 const (
 	DefaultDC = 2
 
@@ -10,14 +8,12 @@ const (
 	handshakeFrameLenKey            = 32
 	handshakeFrameLenIV             = 16
 	handshakeFrameLenConnectionType = 4
-	handshakeFrameLenDC             = 2
 
 	handshakeFrameOffsetStart          = 8
 	handshakeFrameOffsetKey            = handshakeFrameOffsetStart
 	handshakeFrameOffsetIV             = handshakeFrameOffsetKey + handshakeFrameLenKey
 	handshakeFrameOffsetConnectionType = handshakeFrameOffsetIV + handshakeFrameLenIV
 	handshakeFrameOffsetDC             = handshakeFrameOffsetConnectionType + handshakeFrameLenConnectionType
-	handshakeFrameOffsetEnd            = handshakeFrameOffsetDC + handshakeFrameLenDC
 )
 
 // Connection-Type: Secure. We support only fake tls.
@@ -38,8 +34,7 @@ type handshakeFrame struct {
 }
 
 func (h *handshakeFrame) dc() int {
-	data := h.data[handshakeFrameOffsetDC:handshakeFrameOffsetEnd]
-	idx := int16(binary.LittleEndian.Uint16(data))
+	idx := int16(h.data[handshakeFrameOffsetDC]) | int16(h.data[handshakeFrameOffsetDC+1])<<8 // nolint: gomnd, lll // little endian for int16 is here
 
 	switch {
 	case idx > 0:
diff --git a/mtglib/internal/obfuscated2/handshake_frame_internal_test.go b/mtglib/internal/obfuscated2/handshake_frame_internal_test.go
index 4fedb3f..7b197db 100644
--- a/mtglib/internal/obfuscated2/handshake_frame_internal_test.go
+++ b/mtglib/internal/obfuscated2/handshake_frame_internal_test.go
@@ -1,9 +1,12 @@
 package obfuscated2
 
 import (
+	"crypto/rand"
 	"encoding/base64"
+	"strconv"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/suite"
 )
 
@@ -40,6 +43,30 @@ func (suite *HandshakeFrameTestSuite) TestOk() {
 	suite.EqualValues(2094, inverted.dc())
 }
 
+func (suite *HandshakeFrameTestSuite) TestDC() {
+	testData := map[int16]int{
+		1:  1,
+		-1: 1,
+		0:  DefaultDC,
+	}
+
+	for k, v := range testData {
+		incoming := k
+		expected := v
+
+		suite.T().Run(strconv.Itoa(int(incoming)), func(t *testing.T) {
+			frame := handshakeFrame{}
+
+			rand.Read(frame.data[:]) // nolint: errcheck
+
+			frame.data[handshakeFrameOffsetDC] = byte(incoming)
+			frame.data[handshakeFrameOffsetDC+1] = byte(incoming >> 8)
+
+			assert.Equal(t, expected, frame.dc())
+		})
+	}
+}
+
 func TestHandshakeFrame(t *testing.T) {
 	t.Parallel()
 	suite.Run(t, &HandshakeFrameTestSuite{})