upstream 支持 quic 和 kcp

Author: tursom

.gitignore

@@ -1,6 +1,7 @@
 .vscode
 
-config.json
 config.toml
-/gateway
+/gateway*
+/kcp*
+/quic*
 /logs/

README.md

@@ -10,10 +10,20 @@
 - log
 - KCP 的 data_shards 和 parity_Shards
 - QUIC 的 application_protocols
+- pid_file
+
+### 顶层配置
+
+| 配置     | 类型   | 备注     |
+| -------- | ------ | -------- |
+| pid_file | string | pid 文件 |
+
+> pid_file 在非 windows 平台默认会写入 /var/run/mc-gateway.pid，
+> 在 windows 平台默认不会写入任何文件
 
 ### hosts
 
-hosts 使用期望的 host 做 key，转发的目的地址为 value。参考`config.example.toml`。
+hosts 使用期望的 host 做 key，转发的目的地址为 value。参考`config.example.toml`。默认的 fallback host 配置 key 为 `default`。
 
 ### log
 

cmd/gateway/config.go

@@ -27,6 +27,7 @@ type (
 		Kcp     KcpConfig         `toml:"kcp"`
 		Hosts   map[string]string `toml:"hosts"`
 		Log     LogConfig         `toml:"log"`
+		PidFile string            `toml:"pid_file"`
 	}
 
 	ProtocolConfig struct {
@@ -42,9 +43,9 @@ type (
 	}
 
 	QuicConfig struct {
-		Enable             bool     `toml:"enable"`
-		Port               int      `toml:"port"`
-		ApplicionProtocols []string `toml:"application_protocols"`
+		Enable               bool     `toml:"enable"`
+		Port                 int      `toml:"port"`
+		ApplicationProtocols []string `toml:"application_protocols"`
 	}
 
 	LogConfig struct {
@@ -72,6 +73,8 @@ func loadConfig() error {
 		return err
 	}
 
+	writePIDFile()
+
 	return loadLogger()
 }
 

cmd/gateway/kcp.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"fmt"
+	"net"
 	"sync"
 
 	"github.com/rs/zerolog/log"
@@ -39,3 +40,15 @@ func runKcp(wg *sync.WaitGroup) {
 		go handleRequest(conn)
 	}
 }
+
+func upstreamKcp(host string) net.Conn {
+	conn, err := kcp.DialWithOptions(host, nil, config.Kcp.DataShards, config.Kcp.ParityShards)
+	if err != nil {
+		log.Error().Err(err).
+			Msg("Failed to dial KCP server")
+	}
+	defer conn.Close()
+
+	conn.SetACKNoDelay(true)
+	return conn
+}

cmd/gateway/main.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"strings"
 	"sync"
 	"time"
 
@@ -12,15 +13,15 @@ import (
 )
 
 func main() {
-	if err := writePIDFile(); err != nil {
-		panic(fmt.Sprintf("Failed to write PID file: %v", err))
-	}
-	defer removePIDFile()
-
 	if err := loadConfig(); err != nil {
 		panic(err)
 	}
 
+	if err := writePIDFile(); err != nil {
+		log.Err(err).Msg("Failed to write PID file")
+	}
+	defer removePIDFile()
+
 	watcher := watchConfig()
 	defer watcher.Close()
 
@@ -70,6 +71,7 @@ func runTcp(wg *sync.WaitGroup) {
 			log.Err(err).Msg("Error accepting")
 			continue
 		}
+		setSocketOptions(conn)
 		// 处理连接
 		go handleRequest(conn)
 	}
@@ -96,23 +98,47 @@ func handleRequest(conn net.Conn) {
 	// 确保连接关闭
 	defer conn.Close()
 
-	setSocketOptions(conn)
+	client := mapToHost(conn)
+	if client == nil {
+		return
+	}
+	defer client.Close()
 
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	go handleRead(client, conn, &wg)
+	handleWrite(client, conn, nil)
+
+	// 等待所有读写操作完成
+	// 不放在 defer 中，以防报错时无法关闭连接
+	wg.Wait()
+}
+
+func mapToHost(conn net.Conn) net.Conn {
 	buf := make([]byte, 1024)
 	n, err := conn.Read(buf)
 	if err != nil {
 		log.Err(err).
 			Str("client", conn.RemoteAddr().String()).
-			Msg("Error reading hostname")
-		return
+			Msg("failed to reading hostname")
+		return nil
 	}
 	if n == 0 {
 		log.Err(errEmptyBuffer).
 			Str("client", conn.RemoteAddr().String()).
-			Msg("Error: buffer is empty")
-		return
+			Msg("buffer is empty")
+		return nil
 	}
+
 	mc_host := protocol.GetMcHost(buf[:n])
+	if mc_host == "" {
+		log.Err(errEmptyBuffer).
+			Str("client", conn.RemoteAddr().String()).
+			Msg("failed to parse mc host from buffer")
+		return nil
+	}
+
 	host, ok := config.Hosts[mc_host]
 	if !ok {
 		host = config.Hosts["default"]
@@ -122,7 +148,7 @@ func handleRequest(conn net.Conn) {
 			Str("client", conn.RemoteAddr().String()).
 			Str("host", mc_host).
 			Msg("failed to route host")
-		return
+		return nil
 	}
 
 	log.Info().
@@ -131,27 +157,33 @@ func handleRequest(conn net.Conn) {
 		Str("mc", host).
 		Msg("map to host")
 
-	client, err := net.Dial("tcp", host)
-	if err != nil {
-		log.Err(err).Msg("Error dialing")
-		return
+	var client net.Conn
+
+	if host, ok := strings.CutPrefix(host, "quic://"); ok {
+		client = upstreamQuic(host)
+	} else if host, ok := strings.CutPrefix(host, "kcp://"); ok {
+		client = upstreamKcp(host)
+	} else {
+		client = upstreamTcp(host)
+	}
+	if client == nil {
+		return nil
 	}
-	defer client.Close()
-	setSocketOptions(client)
 
 	client.Write(buf[:n])
-	// 不需要 buf 了，释放掉
-	buf = nil
 
-	var wg sync.WaitGroup
-	wg.Add(1)
+	return client
+}
 
-	go handleRead(client, conn, &wg)
-	handleWrite(client, conn, nil)
+func upstreamTcp(host string) net.Conn {
+	conn, err := net.Dial("tcp", host)
+	if err != nil {
+		log.Err(err).Str("host", host).Msg("Error dialing upstream")
+		return nil
+	}
+	setSocketOptions(conn)
+	return conn
 
-	// 等待所有读写操作完成
-	// 不放在 defer 中，以防报错时无法关闭连接
-	wg.Wait()
 }
 
 func handleRead(srv, cli net.Conn, wg *sync.WaitGroup) {

cmd/gateway/pid.go

@@ -0,0 +1,31 @@
+package main
+
+import (
+	"fmt"
+	"os"
+)
+
+var currentPidFile string
+
+func writePIDFile() error {
+	newPidFile := getPidFileFromConfig()
+	if newPidFile == currentPidFile {
+		return nil
+	}
+
+	if currentPidFile != "" {
+		removePIDFile()
+	}
+
+	pid := os.Getpid()
+	if err := os.WriteFile(newPidFile, []byte(fmt.Sprintf("%d\n", pid)), 0644); err != nil {
+		return err
+	}
+
+	currentPidFile = newPidFile
+	return nil
+}
+
+func removePIDFile() {
+	os.Remove(currentPidFile)
+}

cmd/gateway/pid_unix.go

@@ -3,16 +3,10 @@
 
 package main
 
-import (
-	"fmt"
-	"os"
-)
-
-func writePIDFile() error {
-	pid := os.Getpid()
-	return os.WriteFile("/dev/shm/mc-gateway.pid", []byte(fmt.Sprintf("%d\n", pid)), 0644)
-}
-
-func removePIDFile() {
-	os.Remove("/dev/shm/mc-gateway.pid")
+func getPidFileFromConfig() string {
+	pidFile := config.PidFile
+	if pidFile == "" {
+		return "/dev/shm/mc-gateway.pid"
+	}
+	return pidFile
 }

cmd/gateway/pid_windows.go

@@ -3,9 +3,6 @@
 
 package main
 
-func writePIDFile() error {
-	return nil
-}
-
-func removePIDFile() {
+func getPidFileFromConfig() string {
+	return config.PidFile
 }

cmd/gateway/quic.go

@@ -18,10 +18,12 @@ import (
 	"github.com/rs/zerolog/log"
 )
 
-type quicConn struct {
-	quic.Connection
-	quic.Stream
-}
+type (
+	quicConn struct {
+		quic.Connection
+		quic.Stream
+	}
+)
 
 func runQuic(wg *sync.WaitGroup) {
 	if wg != nil {
@@ -59,6 +61,34 @@ func runQuic(wg *sync.WaitGroup) {
 	}
 }
 
+func upstreamQuic(host string) net.Conn {
+	tlsConf := &tls.Config{
+		InsecureSkipVerify: true, // 跳过证书检查
+		NextProtos:         getQuicNextProtos(),
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) // 3s handshake timeout
+	defer cancel()
+
+	conn, err := quic.DialAddr(ctx, host, tlsConf, nil)
+	if err != nil {
+		log.Err(err).Str("host", host).Msg("Failed to dial QUIC")
+		return nil
+	}
+
+	stream, err := conn.OpenStream()
+	if err != nil {
+		log.Err(err).Str("host", host).Msg("Failed to open stream")
+		return nil
+	}
+	log.Info().Str("host", host).Msg("QUIC stream opened")
+
+	return quicConn{
+		Connection: conn,
+		Stream:     stream,
+	}
+}
+
 func handleQuicRequest(conn quic.Connection) {
 	defer conn.CloseWithError(0, "Closing connection")
 
@@ -118,14 +148,24 @@ func generateTLSConfig() (*tls.Config, error) {
 		return nil, err
 	}
 
-	nextProtos := config.Quic.ApplicionProtocols
-	if len(nextProtos) == 0 {
-		nextProtos = []string{"minecraft", "quic", "raw", "h3"} // 默认协议
-	}
-
 	// 返回 tls.Config
 	return &tls.Config{
 		Certificates: []tls.Certificate{cert},
-		NextProtos:   nextProtos,
+		NextProtos:   getQuicNextProtos(),
 	}, nil
 }
+
+func getQuicNextProtos() []string {
+	nextProtos := config.Quic.ApplicationProtocols
+	if len(nextProtos) == 0 {
+		return []string{"minecraft", "quic", "raw", "h3"} // 默认协议
+	}
+	return nextProtos
+}
+
+func (c quicConn) Close() error {
+	if err := c.Stream.Close(); err != nil {
+		log.Err(err).Msg("Failed to close QUIC stream")
+	}
+	return c.Connection.CloseWithError(0, "Closing QUIC connection")
+}