添加 haproxy 协议支持

添加插件功能支持

cmd/gateway/config.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/BurntSushi/toml"
 	"github.com/fsnotify/fsnotify"
+	"github.com/mitchellh/mapstructure"
 	"github.com/rs/zerolog/log"
 )
 
@@ -22,13 +23,14 @@ var (
 
 type (
 	Config struct {
-		Tcp       ProtocolConfig    `toml:"tcp"`
-		Quic      QuicConfig        `toml:"quic"`
-		Kcp       KcpConfig         `toml:"kcp"`
-		WebSocket WebSocketConfig   `toml:"websocket"`
-		Hosts     map[string]string `toml:"hosts"`
-		Log       LogConfig         `toml:"log"`
-		PidFile   string            `toml:"pid_file"`
+		Tcp       ProtocolConfig            `toml:"tcp"`
+		Quic      QuicConfig                `toml:"quic"`
+		Kcp       KcpConfig                 `toml:"kcp"`
+		WebSocket WebSocketConfig           `toml:"websocket"`
+		Hosts     map[string]string         `toml:"hosts"`
+		Log       LogConfig                 `toml:"log"`
+		PidFile   string                    `toml:"pid_file"`
+		Plugin    map[string]map[string]any `toml:"plugin"`
 	}
 
 	ProtocolConfig struct {
@@ -115,7 +117,13 @@ func loadConfig() error {
 
 	writePIDFile()
 
-	return loadLogger()
+	if err := loadLogger(); err != nil {
+		return err
+	}
+
+	loadPlugins()
+
+	return nil
 }
 
 func watchConfig() *fsnotify.Watcher {
@@ -164,3 +172,23 @@ func watchConfig() *fsnotify.Watcher {
 
 	return watcher
 }
+
+func loadPluginConfig(cfg map[string]any, pluginCfg any) error {
+	log.Info().
+		Any("config", cfg).
+		Msg("Loading plugin config")
+
+	decoder, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{
+		Result:  pluginCfg,
+		TagName: "toml",
+	})
+	if err != nil {
+		return err
+	}
+
+	if err := decoder.Decode(cfg); err != nil {
+		return err
+	}
+
+	return nil
+}

cmd/gateway/haproxy.go

@@ -0,0 +1,52 @@
+package main
+
+import (
+	"net"
+
+	proxyproto "github.com/pires/go-proxyproto"
+	"github.com/rs/zerolog/log"
+)
+
+func haProxyUpstream(source net.Conn, host string) net.Conn {
+	target, err := net.ResolveTCPAddr("tcp", host)
+	if err != nil {
+		log.Err(err).Msg("failed to resolve TCP address")
+		return nil
+	}
+
+	conn, err := net.DialTCP("tcp", nil, target)
+	if err != nil {
+		log.Err(err).Msg("failed to dial TCP")
+		return nil
+	}
+
+	sourceAddr, err := net.ResolveTCPAddr(
+		source.RemoteAddr().Network(),
+		source.RemoteAddr().String(),
+	)
+	if err != nil {
+		log.Err(err).Msg("failed to resolve TCP address")
+		return nil
+	}
+
+	TransportProtocol := proxyproto.TCPv4
+	if sourceAddr.IP.To4() == nil {
+		TransportProtocol = proxyproto.TCPv6
+	}
+
+	header := &proxyproto.Header{
+		Version:           1,
+		Command:           proxyproto.PROXY,
+		TransportProtocol: TransportProtocol,
+		SourceAddr:        sourceAddr,
+		DestinationAddr:   target,
+	}
+	// After the connection was created write the proxy headers first
+	_, err = header.WriteTo(conn)
+	if err != nil {
+		log.Err(err).Msg("failed to write proxy header")
+		return nil
+	}
+
+	return conn
+}

cmd/gateway/main.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
+	"github.com/tursom/mc-gateway/plugin/api"
 	"github.com/tursom/mc-gateway/protocol"
 )
 
@@ -26,16 +27,15 @@ func main() {
 
 	go handleLogRotate()
 
-	var wg sync.WaitGroup
-	defer wg.Wait()
+	defer exitWaitGroup.Wait()
 
 	for _, service := range services {
 		if !*service.enable {
 			continue
 		}
 
-		wg.Add(1)
-		go service.run(&wg)
+		exitWaitGroup.Add(1)
+		go service.run(&exitWaitGroup)
 	}
 }
 
@@ -121,12 +121,26 @@ func mapToHost(conn net.Conn) net.Conn {
 
 	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)
+	ok, err = invokeFirstHookHandler(api.HookUpstream, Handler2[net.Conn, string, bool](conn, host), func(handler func(net.Conn, string) (net.Conn, error)) error {
+		var err error
+		client, err = handler(conn, host)
+		return err
+	})
+	if err != nil {
+		log.Err(err).Msg("Failed to invoke upstream hook")
+		return nil
+	}
+
+	if !ok {
+		if host, ok := strings.CutPrefix(host, "quic://"); ok {
+			client = upstreamQuic(host)
+		} else if host, ok := strings.CutPrefix(host, "kcp://"); ok {
+			client = upstreamKcp(host)
+		} else if host, ok := strings.CutPrefix(host, "haproxy://"); ok {
+			client = haProxyUpstream(conn, host)
+		} else {
+			client = upstreamTcp(host)
+		}
 	}
 	if client == nil {
 		return nil

cmd/gateway/plugin.go

@@ -0,0 +1,175 @@
+package main
+
+import (
+	"net"
+	"plugin"
+	"sync"
+
+	"github.com/rs/zerolog/log"
+	"github.com/tursom/mc-gateway/plugin/api"
+)
+
+var (
+	exitWaitGroup sync.WaitGroup
+
+	pluginLock sync.RWMutex
+	plugins    = make(map[string]api.Plugin)
+
+	hooks = make(map[string]map[string]any)
+)
+
+type (
+	Gateway struct {
+		pluginId string
+	}
+)
+
+func loadPlugins() {
+	pluginLock.Lock()
+	defer pluginLock.Unlock()
+
+	for pluginKey, pluginConfig := range config.Plugin {
+		if enable, ok := pluginConfig["enable"].(bool); ok && enable {
+			pluginFile := pluginKey
+			if file, ok := pluginConfig["file"].(string); ok {
+				pluginFile = file
+			}
+
+			if _, ok := plugins[pluginKey]; ok {
+				continue
+			}
+
+			gateway := &Gateway{
+				pluginId: pluginKey,
+			}
+
+			log.Info().Str("plugin", pluginKey).Msg("Loading plugin")
+			p, err := plugin.Open(pluginFile + ".so")
+			if err != nil {
+				log.Err(err).Str("plugin", pluginKey).Msg("Failed to open plugin")
+				continue
+			}
+
+			pluginSymbol, err := p.Lookup("Plugin")
+			if err != nil {
+				log.Err(err).Str("plugin", pluginKey).Msg("Failed to lookup plugin")
+				continue
+			}
+			pluginFactory, ok := pluginSymbol.(func() api.Plugin)
+			if !ok {
+				log.Err(err).Str("plugin", pluginKey).Msg("Invalid plugin factory signature")
+				continue
+			}
+			pluginInstance := pluginFactory()
+			cfgObj := pluginInstance.NewConfigObj()
+			if err := loadPluginConfig(pluginConfig, cfgObj); err != nil {
+				log.Err(err).Str("plugin", pluginKey).Msg("Failed to load plugin config")
+				continue
+			}
+
+			if err := pluginInstance.ReloadConfig(cfgObj); err != nil {
+				log.Err(err).Str("plugin", pluginKey).Msg("Failed to reload plugin config")
+				continue
+			}
+
+			if err := pluginInstance.Init(gateway); err != nil {
+				log.Err(err).Str("plugin", pluginKey).Msg("Failed to initialize plugin")
+				continue
+			}
+
+			plugins[pluginKey] = pluginInstance
+			hooks[pluginKey] = make(map[string]any)
+			log.Info().Str("plugin", pluginKey).Msg("Plugin loaded successfully")
+		} else {
+			if plugin, ok := plugins[pluginKey]; ok {
+				if err := plugin.Destroy(); err != nil {
+					log.Err(err).Str("plugin", pluginKey).Msg("Failed to destroy plugin")
+				}
+				log.Info().Str("plugin", pluginKey).Msg("Plugin disabled")
+			}
+			delete(plugins, pluginKey)
+			delete(hooks, pluginKey)
+		}
+	}
+}
+
+// HandleConn implements api.Gateway.
+func (g *Gateway) HandleConn(conn net.Conn) {
+	go handleRequest(conn)
+}
+
+// Hook implements api.Gateway.
+func (g *Gateway) Hook(hook string, handler any) error {
+	pluginLock.Lock()
+	defer pluginLock.Unlock()
+
+	hooks[g.pluginId][hook] = handler
+	return nil
+}
+
+// ExitWaitGroup implements api.Gateway.
+func (g *Gateway) ExitWaitGroup() *sync.WaitGroup {
+	return &exitWaitGroup
+}
+
+// TestOp implements api.Gateway.
+func (g *Gateway) TestOp() {
+	panic("unimplemented")
+}
+
+func Handler1[T1, R any](t1 T1) func(func(T1) R) R {
+	return func(acceptor func(T1) R) R {
+		return acceptor(t1)
+	}
+}
+
+func Handler2[T1, T2, R any](t1 T1, t2 T2) func(func(T1, T2) R) R {
+	return func(acceptor func(T1, T2) R) R {
+		return acceptor(t1, t2)
+	}
+}
+
+func Handler1R2[T1, R1, R2 any](t1 T1) func(func(T1) (R1, R2)) (R1, R2) {
+	return func(acceptor func(T1) (R1, R2)) (R1, R2) {
+		return acceptor(t1)
+	}
+}
+
+func invokeFirstHookHandler[Acceptor, Handler any](
+	hook api.HookType[Acceptor, Handler],
+	acceptor func(Acceptor) bool,
+	handelr func(Handler) error,
+) (bool, error) {
+	pluginLock.RLock()
+	defer pluginLock.RUnlock()
+
+	for _, handlers := range hooks {
+		if handler, ok := handlers[hook.Key()].(api.HookHandler[Acceptor, Handler]); ok && acceptor(handler.Acceptor()) {
+			if err := handelr(handler.Handler()); err != nil {
+				return true, err
+			}
+			return true, nil
+		}
+	}
+
+	return false, nil
+}
+
+func invokeAllHookHandler[Acceptor, Handler any](
+	hook api.HookType[Acceptor, Handler],
+	acceptor func(Acceptor) bool,
+	handelr func(Handler) error,
+) error {
+	pluginLock.RLock()
+	defer pluginLock.RUnlock()
+
+	for _, handlers := range hooks {
+		if handler, ok := handlers[hook.Key()].(api.HookHandler[Acceptor, Handler]); ok && acceptor(handler.Acceptor()) {
+			if err := handelr(handler.Handler()); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}