From 32d34c98692d6ff49d39ee88bf0a5e2525f74a69 Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Thu, 5 Jul 2018 14:47:46 +0200
Subject: [PATCH] fix dc id selection

---
 proxy/mtproto/auth.go      | 11 ++++++++++-
 proxy/mtproto/auth_test.go | 15 +++++++++++++++
 proxy/mtproto/server.go    |  7 +++++--
 3 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/proxy/mtproto/auth.go b/proxy/mtproto/auth.go
index 962bc1c39..973561cac 100644
--- a/proxy/mtproto/auth.go
+++ b/proxy/mtproto/auth.go
@@ -22,7 +22,11 @@ type Authentication struct {
 }
 
 func (a *Authentication) DataCenterID() uint16 {
-	return ((uint16(a.Header[61]) << 8) | uint16(a.Header[60])) % uint16(len(dcList))
+	x := ((int16(a.Header[61]) << 8) | int16(a.Header[60]))
+	if x < 0 {
+		x = -x
+	}
+	return uint16(x) - 1
 }
 
 func (a *Authentication) ApplySecret(b []byte) {
@@ -47,6 +51,11 @@ func generateRandomBytes(random []byte) {
 			continue
 		}
 
+		random[56] = 0xef
+		random[57] = 0xef
+		random[58] = 0xef
+		random[59] = 0xef
+
 		return
 	}
 }
diff --git a/proxy/mtproto/auth_test.go b/proxy/mtproto/auth_test.go
index 82630774c..e2778d0c2 100644
--- a/proxy/mtproto/auth_test.go
+++ b/proxy/mtproto/auth_test.go
@@ -1,6 +1,7 @@
 package mtproto_test
 
 import (
+	"bytes"
 	"crypto/rand"
 	"testing"
 
@@ -21,3 +22,17 @@ func TestInverse(t *testing.T) {
 	bii := Inverse(bi)
 	assert(bii, Equals, b)
 }
+
+func TestAuthenticationReadWrite(t *testing.T) {
+	assert := With(t)
+
+	a := NewAuthentication()
+	b := bytes.NewReader(a.Header[:])
+	a2, err := ReadAuthentication(b)
+	assert(err, IsNil)
+
+	assert(a.EncodingKey[:], Equals, a2.DecodingKey[:])
+	assert(a.EncodingNonce[:], Equals, a2.DecodingNonce[:])
+	assert(a.DecodingKey[:], Equals, a2.EncodingKey[:])
+	assert(a.DecodingNonce[:], Equals, a2.EncodingNonce[:])
+}
diff --git a/proxy/mtproto/server.go b/proxy/mtproto/server.go
index a645d470f..8d16225d4 100644
--- a/proxy/mtproto/server.go
+++ b/proxy/mtproto/server.go
@@ -90,6 +90,9 @@ func (s *Server) Process(ctx context.Context, network net.Network, conn internet
 	}
 
 	dcID := auth.DataCenterID()
+	if dcID >= uint16(len(dcList)) {
+		return newError("invalid datacenter id: ", dcID)
+	}
 
 	dest := net.Destination{
 		Network: net.Network_TCP,
@@ -110,7 +113,7 @@ func (s *Server) Process(ctx context.Context, network net.Network, conn internet
 		defer timer.SetTimeout(sPolicy.Timeouts.DownlinkOnly)
 
 		reader := buf.NewReader(crypto.NewCryptionReader(decryptor, conn))
-		return buf.Copy(reader, link.Writer)
+		return buf.Copy(reader, link.Writer, buf.UpdateActivity(timer))
 	}
 
 	response := func() error {
@@ -118,7 +121,7 @@ func (s *Server) Process(ctx context.Context, network net.Network, conn internet
 
 		encryptor := crypto.NewAesCTRStream(auth.EncodingKey[:], auth.EncodingNonce[:])
 		writer := buf.NewWriter(crypto.NewCryptionWriter(encryptor, conn))
-		return buf.Copy(link.Reader, writer)
+		return buf.Copy(link.Reader, writer, buf.UpdateActivity(timer))
 	}
 
 	var responseDoneAndCloseWriter = task.Single(response, task.OnSuccess(task.Close(link.Writer)))