v1.0.0

proxy/mtproto/auth.go

@@ -0,0 +1,150 @@
+package mtproto
+
+import (
+	"context"
+	"crypto/rand"
+	"crypto/sha256"
+	"io"
+	"sync"
+
+	"github.com/xtls/xray-core/v1/common"
+)
+
+const (
+	HeaderSize = 64
+)
+
+type SessionContext struct {
+	ConnectionType [4]byte
+	DataCenterID   uint16
+}
+
+func DefaultSessionContext() SessionContext {
+	return SessionContext{
+		ConnectionType: [4]byte{0xef, 0xef, 0xef, 0xef},
+		DataCenterID:   0,
+	}
+}
+
+type contextKey int32
+
+const (
+	sessionContextKey contextKey = iota
+)
+
+func ContextWithSessionContext(ctx context.Context, c SessionContext) context.Context {
+	return context.WithValue(ctx, sessionContextKey, c)
+}
+
+func SessionContextFromContext(ctx context.Context) SessionContext {
+	if c := ctx.Value(sessionContextKey); c != nil {
+		return c.(SessionContext)
+	}
+	return DefaultSessionContext()
+}
+
+type Authentication struct {
+	Header        [HeaderSize]byte
+	DecodingKey   [32]byte
+	EncodingKey   [32]byte
+	DecodingNonce [16]byte
+	EncodingNonce [16]byte
+}
+
+func (a *Authentication) DataCenterID() uint16 {
+	x := ((int16(a.Header[61]) << 8) | int16(a.Header[60]))
+	if x < 0 {
+		x = -x
+	}
+	return uint16(x) - 1
+}
+
+func (a *Authentication) ConnectionType() [4]byte {
+	var x [4]byte
+	copy(x[:], a.Header[56:60])
+	return x
+}
+
+func (a *Authentication) ApplySecret(b []byte) {
+	a.DecodingKey = sha256.Sum256(append(a.DecodingKey[:], b...))
+	a.EncodingKey = sha256.Sum256(append(a.EncodingKey[:], b...))
+}
+
+func generateRandomBytes(random []byte, connType [4]byte) {
+	for {
+		common.Must2(rand.Read(random))
+
+		if random[0] == 0xef {
+			continue
+		}
+
+		val := (uint32(random[3]) << 24) | (uint32(random[2]) << 16) | (uint32(random[1]) << 8) | uint32(random[0])
+		if val == 0x44414548 || val == 0x54534f50 || val == 0x20544547 || val == 0x4954504f || val == 0xeeeeeeee {
+			continue
+		}
+
+		if (uint32(random[7])<<24)|(uint32(random[6])<<16)|(uint32(random[5])<<8)|uint32(random[4]) == 0x00000000 {
+			continue
+		}
+
+		copy(random[56:60], connType[:])
+
+		return
+	}
+}
+
+func NewAuthentication(sc SessionContext) *Authentication {
+	auth := getAuthenticationObject()
+	random := auth.Header[:]
+	generateRandomBytes(random, sc.ConnectionType)
+	copy(auth.EncodingKey[:], random[8:])
+	copy(auth.EncodingNonce[:], random[8+32:])
+	keyivInverse := Inverse(random[8 : 8+32+16])
+	copy(auth.DecodingKey[:], keyivInverse)
+	copy(auth.DecodingNonce[:], keyivInverse[32:])
+	return auth
+}
+
+func ReadAuthentication(reader io.Reader) (*Authentication, error) {
+	auth := getAuthenticationObject()
+
+	if _, err := io.ReadFull(reader, auth.Header[:]); err != nil {
+		putAuthenticationObject(auth)
+		return nil, err
+	}
+
+	copy(auth.DecodingKey[:], auth.Header[8:])
+	copy(auth.DecodingNonce[:], auth.Header[8+32:])
+	keyivInverse := Inverse(auth.Header[8 : 8+32+16])
+	copy(auth.EncodingKey[:], keyivInverse)
+	copy(auth.EncodingNonce[:], keyivInverse[32:])
+
+	return auth, nil
+}
+
+// Inverse returns a new byte array. It is a sequence of bytes when the input is read from end to beginning.Inverse
+// Visible for testing only.
+func Inverse(b []byte) []byte {
+	lenb := len(b)
+	b2 := make([]byte, lenb)
+	for i, v := range b {
+		b2[lenb-i-1] = v
+	}
+	return b2
+}
+
+var (
+	authPool = sync.Pool{
+		New: func() interface{} {
+			return new(Authentication)
+		},
+	}
+)
+
+func getAuthenticationObject() *Authentication {
+	return authPool.Get().(*Authentication)
+}
+
+func putAuthenticationObject(auth *Authentication) {
+	authPool.Put(auth)
+}