update registry

pkg/registry/registry.go

@@ -2,70 +2,115 @@ package registry
 
 import (
 	"errors"
+	"sync"
 
-	"github.com/go-gost/gost/pkg/connector"
-	"github.com/go-gost/gost/pkg/dialer"
-	"github.com/go-gost/gost/pkg/handler"
-	"github.com/go-gost/gost/pkg/listener"
-	"github.com/go-gost/gost/pkg/logger"
+	"github.com/go-gost/gost/pkg/admission"
+	"github.com/go-gost/gost/pkg/auth"
+	"github.com/go-gost/gost/pkg/bypass"
+	"github.com/go-gost/gost/pkg/chain"
+	"github.com/go-gost/gost/pkg/hosts"
+	"github.com/go-gost/gost/pkg/resolver"
+	"github.com/go-gost/gost/pkg/service"
 )
 
 var (
-	ErrDup = errors.New("registry: duplicate instance")
+	ErrDup = errors.New("registry: duplicate object")
 )
 
-type NewListener func(opts ...listener.Option) listener.Listener
-type NewHandler func(opts ...handler.Option) handler.Handler
-type NewDialer func(opts ...dialer.Option) dialer.Dialer
-type NewConnector func(opts ...connector.Option) connector.Connector
-
 var (
-	listeners  = make(map[string]NewListener)
-	handlers   = make(map[string]NewHandler)
-	dialers    = make(map[string]NewDialer)
-	connectors = make(map[string]NewConnector)
+	listenerReg  Registry[NewListener]  = &listenerRegistry{}
+	handlerReg   Registry[NewHandler]   = &handlerRegistry{}
+	dialerReg    Registry[NewDialer]    = &dialerRegistry{}
+	connectorReg Registry[NewConnector] = &connectorRegistry{}
+
+	serviceReg   Registry[service.Service]     = &serviceRegistry{}
+	chainReg     Registry[chain.Chainer]       = &chainRegistry{}
+	autherReg    Registry[auth.Authenticator]  = &autherRegistry{}
+	admissionReg Registry[admission.Admission] = &admissionRegistry{}
+	bypassReg    Registry[bypass.Bypass]       = &bypassRegistry{}
+	resolverReg  Registry[resolver.Resolver]   = &resolverRegistry{}
+	hostsReg     Registry[hosts.HostMapper]    = &hostsRegistry{}
 )
 
-func RegisterListener(name string, newf NewListener) {
-	if listeners[name] != nil {
-		logger.Default().Fatalf("register duplicate listener: %s", name)
+type Registry[T any] interface {
+	Register(name string, v T) error
+	Unregister(name string)
+	IsRegistered(name string) bool
+	Get(name string) T
+}
+
+type registry struct {
+	m sync.Map
+}
+
+func (r *registry) Register(name string, v any) error {
+	if name == "" || v == nil {
+		return nil
 	}
-	listeners[name] = newf
-}
-
-func GetListener(name string) NewListener {
-	return listeners[name]
-}
-
-func RegisterHandler(name string, newf NewHandler) {
-	if handlers[name] != nil {
-		logger.Default().Fatalf("register duplicate handler: %s", name)
+	if _, loaded := r.m.LoadOrStore(name, v); loaded {
+		return ErrDup
 	}
-	handlers[name] = newf
+
+	return nil
 }
 
-func GetHandler(name string) NewHandler {
-	return handlers[name]
+func (r *registry) Unregister(name string) {
+	r.m.Delete(name)
 }
 
-func RegisterDialer(name string, newf NewDialer) {
-	if dialers[name] != nil {
-		logger.Default().Fatalf("register duplicate dialer: %s", name)
+func (r *registry) IsRegistered(name string) bool {
+	_, ok := r.m.Load(name)
+	return ok
+}
+
+func (r *registry) Get(name string) any {
+	if name == "" {
+		return nil
 	}
-	dialers[name] = newf
+	v, _ := r.m.Load(name)
+	return v
 }
 
-func GetDialer(name string) NewDialer {
-	return dialers[name]
+func ListenerRegistry() Registry[NewListener] {
+	return listenerReg
 }
 
-func RegiserConnector(name string, newf NewConnector) {
-	if connectors[name] != nil {
-		logger.Default().Fatalf("register duplicate connector: %s", name)
-	}
-	connectors[name] = newf
+func HandlerRegistry() Registry[NewHandler] {
+	return handlerReg
 }
 
-func GetConnector(name string) NewConnector {
-	return connectors[name]
+func DialerRegistry() Registry[NewDialer] {
+	return dialerReg
+}
+
+func ConnectorRegistry() Registry[NewConnector] {
+	return connectorReg
+}
+
+func ServiceRegistry() Registry[service.Service] {
+	return serviceReg
+}
+
+func ChainRegistry() Registry[chain.Chainer] {
+	return chainReg
+}
+
+func AutherRegistry() Registry[auth.Authenticator] {
+	return autherReg
+}
+
+func AdmissionRegistry() Registry[admission.Admission] {
+	return admissionReg
+}
+
+func BypassRegistry() Registry[bypass.Bypass] {
+	return bypassReg
+}
+
+func ResolverRegistry() Registry[resolver.Resolver] {
+	return resolverReg
+}
+
+func HostsRegistry() Registry[hosts.HostMapper] {
+	return hostsReg
 }