chore: improve code

nxtreaming · Jun 5, 2024 · b84638a · b84638a
1 parent 4499c72
commit b84638a
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 13 deletions.
diff --git a/internal/relay/relay.go b/internal/relay/relay.go
@@ -19,12 +19,12 @@ func (r *Relay) UniqueID() string {
 	return r.cfg.Label
 }
 
-func NewRelay(cfg *conf.Config, connMgr cmgr.Cmgr) (*Relay, error) {
-	base := transporter.NewBaseTransporter(cfg, connMgr)
-	s, err := transporter.NewRelayServer(cfg.ListenType, base)
+func NewRelay(cfg *conf.Config, cmgr cmgr.Cmgr) (*Relay, error) {
+	s, err := transporter.NewRelayServer(cfg, cmgr)
 	if err != nil {
 		return nil, err
 	}
+
 	r := &Relay{
 		relayServer: s,
 		cfg:         cfg,

diff --git a/internal/transporter/base.go b/internal/transporter/base.go
@@ -19,10 +19,11 @@ import (
 )
 
 type baseTransporter struct {
+	cfg *conf.Config
+	l   *zap.SugaredLogger
+
 	cmgr       cmgr.Cmgr
-	cfg        *conf.Config
 	tCPRemotes lb.RoundRobin
-	l          *zap.SugaredLogger
 }
 
 func NewBaseTransporter(cfg *conf.Config, cmgr cmgr.Cmgr) *baseTransporter {

diff --git a/internal/transporter/interface.go b/internal/transporter/interface.go
@@ -3,7 +3,9 @@ package transporter
 import (
 	"net"
 
+	"github.com/Ehco1996/ehco/internal/cmgr"
 	"github.com/Ehco1996/ehco/internal/constant"
+	"github.com/Ehco1996/ehco/internal/relay/conf"
 	"github.com/Ehco1996/ehco/pkg/lb"
 )
 
@@ -14,8 +16,8 @@ type RelayClient interface {
 	RelayTCPConn(c net.Conn, handshakeF TCPHandShakeF) error
 }
 
-func NewRelayClient(relayType string, base *baseTransporter) (RelayClient, error) {
-	switch relayType {
+func newRelayClient(base *baseTransporter) (RelayClient, error) {
+	switch base.cfg.TransportType {
 	case constant.RelayTypeRaw:
 		return newRawClient(base)
 	case constant.RelayTypeWS:
@@ -27,7 +29,7 @@ func NewRelayClient(relayType string, base *baseTransporter) (RelayClient, error
 	case constant.RelayTypeMTCP:
 		return newMtcpClient(base)
 	default:
-		panic("unsupported transport type")
+		panic("unsupported transport type" + base.cfg.TransportType)
 	}
 }
 
@@ -36,8 +38,9 @@ type RelayServer interface {
 	Close() error
 }
 
-func NewRelayServer(relayType string, base *baseTransporter) (RelayServer, error) {
-	switch relayType {
+func NewRelayServer(cfg *conf.Config, cmgr cmgr.Cmgr) (RelayServer, error) {
+	base := NewBaseTransporter(cfg, cmgr)
+	switch cfg.ListenType {
 	case constant.RelayTypeRaw:
 		return newRawServer(base)
 	case constant.RelayTypeWS:
@@ -49,6 +52,6 @@ func NewRelayServer(relayType string, base *baseTransporter) (RelayServer, error
 	case constant.RelayTypeMTCP:
 		return newMtcpServer(base)
 	default:
-		panic("unsupported transport type")
+		panic("unsupported transport type" + cfg.ListenType)
 	}
 }
diff --git a/internal/transporter/raw.go b/internal/transporter/raw.go
@@ -57,7 +57,7 @@ func newRawServer(base *baseTransporter) (*RawServer, error) {
 	if err != nil {
 		return nil, err
 	}
-	relayer, err := NewRelayClient(base.cfg.TransportType, base)
+	relayer, err := newRelayClient(base)
 	if err != nil {
 		return nil, err
 	}

diff --git a/internal/transporter/ws.go b/internal/transporter/ws.go
@@ -71,7 +71,7 @@ func newWsServer(base *baseTransporter) (*WsServer, error) {
 	e.GET("/", echo.WrapHandler(web.MakeIndexF()))
 	e.GET("/handshake/", echo.WrapHandler(http.HandlerFunc(s.HandleRequest)))
 	s.e = e
-	relayer, err := NewRelayClient(base.cfg.TransportType, base)
+	relayer, err := newRelayClient(base)
 	if err != nil {
 		return nil, err
 	}