From bcc7b78ff784adf1f2d964fcd30aa7d8cedb271b Mon Sep 17 00:00:00 2001
From: RPRX <63339210+rprx@users.noreply.github.com>
Date: Tue, 6 Oct 2020 16:25:02 +0000
Subject: [PATCH] Add XTLS support to mKCP (#267)

---
 infra/conf/transport_internet.go   |  4 ++--
 transport/internet/kcp/dialer.go   | 11 +++++-----
 transport/internet/kcp/listener.go | 34 ++++++++++++++++++------------
 3 files changed, 29 insertions(+), 20 deletions(-)

diff --git a/infra/conf/transport_internet.go b/infra/conf/transport_internet.go
index 65f4151b6..3906badda 100644
--- a/infra/conf/transport_internet.go
+++ b/infra/conf/transport_internet.go
@@ -496,8 +496,8 @@ func (c *StreamConfig) Build() (*internet.StreamConfig, error) {
 		config.SecurityType = tm.Type
 	}
 	if strings.EqualFold(c.Security, "xtls") {
-		if config.ProtocolName != "tcp" && config.ProtocolName != "domainsocket" {
-			return nil, newError("XTLS only supports TCP and DomainSocket for now.")
+		if config.ProtocolName != "tcp" && config.ProtocolName != "mkcp" && config.ProtocolName != "domainsocket" {
+			return nil, newError("XTLS only supports TCP, mKCP and DomainSocket for now.")
 		}
 		xtlsSettings := c.XTLSSettings
 		if xtlsSettings == nil {
diff --git a/transport/internet/kcp/dialer.go b/transport/internet/kcp/dialer.go
index b0531b500..bc481ae3d 100644
--- a/transport/internet/kcp/dialer.go
+++ b/transport/internet/kcp/dialer.go
@@ -4,7 +4,6 @@ package kcp
 
 import (
 	"context"
-	"crypto/tls"
 	"io"
 	"sync/atomic"
 
@@ -13,7 +12,8 @@ import (
 	"v2ray.com/core/common/dice"
 	"v2ray.com/core/common/net"
 	"v2ray.com/core/transport/internet"
-	v2tls "v2ray.com/core/transport/internet/tls"
+	"v2ray.com/core/transport/internet/tls"
+	"v2ray.com/core/transport/internet/xtls"
 )
 
 var (
@@ -88,9 +88,10 @@ func DialKCP(ctx context.Context, dest net.Destination, streamSettings *internet
 
 	var iConn internet.Connection = session
 
-	if config := v2tls.ConfigFromStreamSettings(streamSettings); config != nil {
-		tlsConn := tls.Client(iConn, config.GetTLSConfig(v2tls.WithDestination(dest)))
-		iConn = tlsConn
+	if config := tls.ConfigFromStreamSettings(streamSettings); config != nil {
+		iConn = tls.Client(iConn, config.GetTLSConfig(tls.WithDestination(dest)))
+	} else if config := xtls.ConfigFromStreamSettings(streamSettings); config != nil {
+		iConn = xtls.Client(iConn, config.GetXTLSConfig(xtls.WithDestination(dest)))
 	}
 
 	return iConn, nil
diff --git a/transport/internet/kcp/listener.go b/transport/internet/kcp/listener.go
index 8e141c6b5..e47ff84a7 100644
--- a/transport/internet/kcp/listener.go
+++ b/transport/internet/kcp/listener.go
@@ -5,15 +5,18 @@ package kcp
 import (
 	"context"
 	"crypto/cipher"
-	"crypto/tls"
+	gotls "crypto/tls"
 	"sync"
 
+	goxtls "github.com/xtls/go"
+
 	"v2ray.com/core/common"
 	"v2ray.com/core/common/buf"
 	"v2ray.com/core/common/net"
 	"v2ray.com/core/transport/internet"
-	v2tls "v2ray.com/core/transport/internet/tls"
+	"v2ray.com/core/transport/internet/tls"
 	"v2ray.com/core/transport/internet/udp"
+	"v2ray.com/core/transport/internet/xtls"
 )
 
 type ConnectionID struct {
@@ -25,14 +28,15 @@ type ConnectionID struct {
 // Listener defines a server listening for connections
 type Listener struct {
 	sync.Mutex
-	sessions  map[ConnectionID]*Connection
-	hub       *udp.Hub
-	tlsConfig *tls.Config
-	config    *Config
-	reader    PacketReader
-	header    internet.PacketHeader
-	security  cipher.AEAD
-	addConn   internet.ConnHandler
+	sessions   map[ConnectionID]*Connection
+	hub        *udp.Hub
+	tlsConfig  *gotls.Config
+	xtlsConfig *goxtls.Config
+	config     *Config
+	reader     PacketReader
+	header     internet.PacketHeader
+	security   cipher.AEAD
+	addConn    internet.ConnHandler
 }
 
 func NewListener(ctx context.Context, address net.Address, port net.Port, streamSettings *internet.MemoryStreamConfig, addConn internet.ConnHandler) (*Listener, error) {
@@ -57,9 +61,12 @@ func NewListener(ctx context.Context, address net.Address, port net.Port, stream
 		addConn:  addConn,
 	}
 
-	if config := v2tls.ConfigFromStreamSettings(streamSettings); config != nil {
+	if config := tls.ConfigFromStreamSettings(streamSettings); config != nil {
 		l.tlsConfig = config.GetTLSConfig()
 	}
+	if config := xtls.ConfigFromStreamSettings(streamSettings); config != nil {
+		l.xtlsConfig = config.GetXTLSConfig()
+	}
 
 	hub, err := udp.ListenUDP(ctx, address, port, streamSettings, udp.HubCapacity(1024))
 	if err != nil {
@@ -131,8 +138,9 @@ func (l *Listener) OnReceive(payload *buf.Buffer, src net.Destination) {
 		}, writer, l.config)
 		var netConn internet.Connection = conn
 		if l.tlsConfig != nil {
-			tlsConn := tls.Server(conn, l.tlsConfig)
-			netConn = tlsConn
+			netConn = gotls.Server(conn, l.tlsConfig)
+		} else if l.xtlsConfig != nil {
+			netConn = goxtls.Server(conn, l.xtlsConfig)
 		}
 
 		l.addConn(netConn)