Skip to content

Commit

Permalink
fix dc id selection
Browse files Browse the repository at this point in the history
  • Loading branch information
DarienRaymond committed Jul 5, 2018
1 parent b16d2e9 commit 32d34c9
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 3 deletions.
11 changes: 10 additions & 1 deletion proxy/mtproto/auth.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -47,6 +51,11 @@ func generateRandomBytes(random []byte) {
continue
}

random[56] = 0xef
random[57] = 0xef
random[58] = 0xef
random[59] = 0xef

return
}
}
Expand Down
15 changes: 15 additions & 0 deletions proxy/mtproto/auth_test.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package mtproto_test

import (
"bytes"
"crypto/rand"
"testing"

Expand All @@ -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[:])
}
7 changes: 5 additions & 2 deletions proxy/mtproto/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -110,15 +113,15 @@ 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 {
defer timer.SetTimeout(sPolicy.Timeouts.UplinkOnly)

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)))
Expand Down

0 comments on commit 32d34c9

Please sign in to comment.