refactor: configs

haveachin · Feb 10, 2024 · e06c332 · e06c332
1 parent e401591
commit e06c332
Show file tree

Hide file tree

Showing 6 changed files with 199 additions and 131 deletions.
diff --git a/pkg/infrared/filter.go b/pkg/infrared/filter.go
@@ -14,29 +14,24 @@ func (f FilterFunc) Filter(c net.Conn) error {
 	return f(c)
 }
 
-type FilterConfigFunc func(cfg *FiltersConfig)
-
-func WithFilterConfig(c FiltersConfig) FilterConfigFunc {
-	return func(cfg *FiltersConfig) {
-		*cfg = c
-	}
-}
-
 type FiltersConfig struct {
 	RateLimiter *RateLimiterConfig `yaml:"rateLimiter"`
 }
 
+func NewFilterConfig() FiltersConfig {
+	rl := NewRateLimiterConfig()
+
+	return FiltersConfig{
+		RateLimiter: &rl,
+	}
+}
+
 type Filter struct {
 	cfg       FiltersConfig
 	filterers []Filterer
 }
 
-func NewFilter(fns ...FilterConfigFunc) Filter {
-	var cfg FiltersConfig
-	for _, fn := range fns {
-		fn(&cfg)
-	}
-
+func NewFilter(cfg FiltersConfig) Filter {
 	filterers := make([]Filterer, 0)
 
 	if cfg.RateLimiter != nil {

diff --git a/pkg/infrared/infrared.go b/pkg/infrared/infrared.go
@@ -5,7 +5,6 @@ import (
 	"io"
 	"net"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/haveachin/infrared/pkg/infrared/protocol"
@@ -22,17 +21,11 @@ type Config struct {
 
 func NewConfig() Config {
 	return Config{
-		BindAddr:         ":25565",
-		KeepAliveTimeout: 30 * time.Second,
-		FiltersConfig: FiltersConfig{
-			RateLimiter: &RateLimiterConfig{
-				RequestLimit: 10,
-				WindowLength: time.Second,
-			},
-		},
-		ProxyProtocolConfig: ProxyProtocolConfig{
-			TrustedCIDRs: make([]string, 0),
-		},
+		BindAddr:            ":25565",
+		KeepAliveTimeout:    30 * time.Second,
+		ServerConfigs:       make([]ServerConfig, 0),
+		FiltersConfig:       NewFilterConfig(),
+		ProxyProtocolConfig: NewProxyProtocolConfig(),
 	}
 }
 
@@ -41,37 +34,28 @@ func (cfg Config) WithBindAddr(bindAddr string) Config {
 	return cfg
 }
 
-func (cfg Config) AddServerConfig(fns ...ServerConfigFunc) Config {
-	var sCfg ServerConfig
-	for _, fn := range fns {
-		fn(&sCfg)
-	}
-	cfg.ServerConfigs = append(cfg.ServerConfigs, sCfg)
-	return cfg
-}
-
 func (cfg Config) WithKeepAliveTimeout(d time.Duration) Config {
 	cfg.KeepAliveTimeout = d
 	return cfg
 }
 
-func (cfg Config) WithProxyProtocolReceive(receive bool) Config {
-	cfg.ProxyProtocolConfig.Receive = receive
+func (cfg Config) WithServerConfigs(sCfgs ...ServerConfig) Config {
+	cfg.ServerConfigs = sCfgs
 	return cfg
 }
 
-func (cfg Config) WithProxyProtocolTrustedCIDRs(trustedCIDRs ...string) Config {
-	cfg.ProxyProtocolConfig.TrustedCIDRs = trustedCIDRs
+func (cfg Config) AddServerConfigs(sCfgs ...ServerConfig) Config {
+	cfg.ServerConfigs = append(cfg.ServerConfigs, sCfgs...)
 	return cfg
 }
 
-func (cfg Config) WithRateLimiterWindowLength(windowLength time.Duration) Config {
-	cfg.FiltersConfig.RateLimiter.WindowLength = windowLength
+func (cfg Config) WithFiltersConfig(fCfg FiltersConfig) Config {
+	cfg.FiltersConfig = fCfg
 	return cfg
 }
 
-func (cfg Config) WithRateLimiterRequestLimit(requestLimit int) Config {
-	cfg.FiltersConfig.RateLimiter.RequestLimit = requestLimit
+func (cfg Config) WithProxyProtocolConfig(ppCfg ProxyProtocolConfig) Config {
+	cfg.ProxyProtocolConfig = ppCfg
 	return cfg
 }
 
@@ -97,38 +81,29 @@ type Infrared struct {
 	Logger                 zerolog.Logger
 	NewListenerFunc        NewListenerFunc
 	NewServerRequesterFunc NewServerRequesterFunc
+	Filters                []Filterer
 
 	cfg Config
 
-	l       net.Listener
-	filter  Filter
-	bufPool sync.Pool
-	conns   map[net.Addr]*clientConn
-	sr      ServerRequester
-}
+	l      net.Listener
+	filter Filter
+	sr     ServerRequester
 
-func New() *Infrared {
-	return NewWithConfig(NewConfig())
+	conns map[net.Addr]*clientConn
 }
 
-func NewWithConfigProvider(prv ConfigProvider) *Infrared {
-	cfg := MustProvideConfig(prv.Config)
-	return NewWithConfig(cfg)
-}
-
-func NewWithConfig(cfg Config) *Infrared {
+func New(cfg Config) *Infrared {
 	return &Infrared{
-		cfg: cfg,
-		bufPool: sync.Pool{
-			New: func() any {
-				b := make([]byte, 1<<15)
-				return &b
-			},
-		},
+		cfg:   cfg,
 		conns: make(map[net.Addr]*clientConn),
 	}
 }
 
+func NewWithConfigProvider(prv ConfigProvider) *Infrared {
+	cfg := MustProvideConfig(prv.Config)
+	return New(cfg)
+}
+
 func (ir *Infrared) initListener() error {
 	ir.Logger.Info().
 		Str("bind", ir.cfg.BindAddr).
@@ -164,7 +139,7 @@ func (ir *Infrared) initListener() error {
 func (ir *Infrared) initServerGateway() error {
 	srvs := make([]*Server, 0)
 	for _, sCfg := range ir.cfg.ServerConfigs {
-		srv, err := NewServer(WithServerConfig(sCfg))
+		srv, err := NewServer(sCfg)
 		if err != nil {
 			return err
 		}
@@ -195,7 +170,8 @@ func (ir *Infrared) init() error {
 		return err
 	}
 
-	ir.filter = NewFilter(WithFilterConfig(ir.cfg.FiltersConfig))
+	ir.filter = NewFilter(ir.cfg.FiltersConfig)
+	ir.filter.filterers = append(ir.filter.filterers, ir.Filters...)
 
 	return nil
 }

diff --git a/pkg/infrared/infrared_test.go b/pkg/infrared/infrared_test.go
@@ -146,7 +146,7 @@ func NewVirtualInfrared(
 	cfg ir.Config,
 	sendProxyProtocol bool,
 ) (*VirtualInfrared, net.Conn) {
-	vir := ir.NewWithConfig(cfg)
+	vir := ir.New(cfg)
 
 	connChan := make(chan net.Conn)
 	vl := &VirtualListener{
@@ -227,8 +227,11 @@ func TestInfrared_SendProxyProtocol_False(t *testing.T) {
 
 func TestInfrared_ReceiveProxyProtocol_True(t *testing.T) {
 	cfg := ir.NewConfig().
-		WithProxyProtocolReceive(true).
-		WithProxyProtocolTrustedCIDRs("127.0.0.1/32")
+		WithProxyProtocolConfig(
+			ir.NewProxyProtocolConfig().
+				WithReceive(true).
+				WithTrustedCIDRs("127.0.0.1/32"),
+		)
 
 	vi, _ := NewVirtualInfrared(cfg, false)
 	go vi.MustListenAndServe(t)
@@ -247,7 +250,10 @@ func TestInfrared_ReceiveProxyProtocol_True(t *testing.T) {
 
 func TestInfrared_ReceiveProxyProtocol_False(t *testing.T) {
 	cfg := ir.NewConfig().
-		WithProxyProtocolReceive(false)
+		WithProxyProtocolConfig(
+			ir.NewProxyProtocolConfig().
+				WithReceive(false),
+		)
 
 	vi, _ := NewVirtualInfrared(cfg, false)
 	go vi.MustListenAndServe(t)
@@ -264,8 +270,10 @@ func TestInfrared_ReceiveProxyProtocol_False(t *testing.T) {
 
 func TestInfrared_ReceiveProxyProtocol_True_ErrNoTrustedCIDRs(t *testing.T) {
 	cfg := ir.NewConfig().
-		WithProxyProtocolReceive(true).
-		WithProxyProtocolTrustedCIDRs()
+		WithProxyProtocolConfig(
+			ir.NewProxyProtocolConfig().
+				WithReceive(true),
+		)
 
 	vi, _ := NewVirtualInfrared(cfg, false)
 
@@ -287,8 +295,11 @@ func TestInfrared_ReceiveProxyProtocol_True_ErrNoTrustedCIDRs(t *testing.T) {
 
 func TestInfrared_ReceiveProxyProtocol_True_UntrustedIP(t *testing.T) {
 	cfg := ir.NewConfig().
-		WithProxyProtocolReceive(true).
-		WithProxyProtocolTrustedCIDRs("127.0.0.1/32")
+		WithProxyProtocolConfig(
+			ir.NewProxyProtocolConfig().
+				WithReceive(true).
+				WithTrustedCIDRs("127.0.0.1/32"),
+		)
 
 	vi, _ := NewVirtualInfrared(cfg, false)
 	go vi.MustListenAndServe(t)

diff --git a/pkg/infrared/proxy_protocol.go b/pkg/infrared/proxy_protocol.go
@@ -17,6 +17,27 @@ type ProxyProtocolConfig struct {
 	TrustedCIDRs []string `yaml:"trustedCIDRs"`
 }
 
+func NewProxyProtocolConfig() ProxyProtocolConfig {
+	return ProxyProtocolConfig{
+		TrustedCIDRs: make([]string, 0),
+	}
+}
+
+func (cfg ProxyProtocolConfig) WithReceive(receive bool) ProxyProtocolConfig {
+	cfg.Receive = receive
+	return cfg
+}
+
+func (cfg ProxyProtocolConfig) WithTrustedCIDRs(trustedCIDRs ...string) ProxyProtocolConfig {
+	cfg.TrustedCIDRs = trustedCIDRs
+	return cfg
+}
+
+func (cfg ProxyProtocolConfig) AddTrustedCIDRs(trustedCIDRs ...string) ProxyProtocolConfig {
+	cfg.TrustedCIDRs = append(cfg.TrustedCIDRs, trustedCIDRs...)
+	return cfg
+}
+
 func newProxyProtocolListener(l net.Listener, trustedCIDRs []string) (net.Listener, error) {
 	if len(trustedCIDRs) == 0 {
 		return nil, ErrNoTrustedCIDRs

diff --git a/pkg/infrared/rate_limiter.go b/pkg/infrared/rate_limiter.go
@@ -17,6 +17,23 @@ type RateLimiterConfig struct {
 	WindowLength time.Duration `yaml:"windowLength"`
 }
 
+func NewRateLimiterConfig() RateLimiterConfig {
+	return RateLimiterConfig{
+		RequestLimit: 10,
+		WindowLength: time.Second,
+	}
+}
+
+func (cfg RateLimiterConfig) WithWindowLength(windowLength time.Duration) RateLimiterConfig {
+	cfg.WindowLength = windowLength
+	return cfg
+}
+
+func (cfg RateLimiterConfig) WithRequestLimit(requestLimit int) RateLimiterConfig {
+	cfg.RequestLimit = requestLimit
+	return cfg
+}
+
 func RateLimit(requestLimit int, windowLength time.Duration, options ...RateLimiterOption) Filterer {
 	return newRateLimiter(requestLimit, windowLength, options...).Filterer()
 }