Darien Raymond 7 роки тому
батько
коміт
32d34c9869
3 змінених файлів з 30 додано та 3 видалено
  1. 10 1
      proxy/mtproto/auth.go
  2. 15 0
      proxy/mtproto/auth_test.go
  3. 5 2
      proxy/mtproto/server.go

+ 10 - 1
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
 	}
 }

+ 15 - 0
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[:])
+}

+ 5 - 2
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)))