diff --git a/pkg/udpproxy/udpproxy.go b/pkg/udpproxy/udpproxy.go index 04df3b8..3b5d052 100644 --- a/pkg/udpproxy/udpproxy.go +++ b/pkg/udpproxy/udpproxy.go @@ -8,6 +8,21 @@ import ( log "github.com/sirupsen/logrus" ) +// Option is a functional option type that allows us to configure the Client. +type Option func(*Client) + +// WithStaleTimeout is a functional option to set the stale session timeout +func WithStaleTimeout(timeout time.Duration) Option { + return func(c *Client) { + if timeout < time.Minute { + log.Warnf("cannot set stale session timeout to less than 1 minute.. defaulting to 5 minutes") + return + } + + c.timeout = timeout + } +} + type Client struct { laddr *net.UDPAddr raddr *net.UDPAddr @@ -16,9 +31,11 @@ type Client struct { mutex sync.RWMutex sessions map[string]*session + + timeout time.Duration } -func New(port, target string) (*Client, error) { +func New(port, target string, options ...Option) (*Client, error) { laddr, err := net.ResolveUDPAddr("udp", port) if err != nil { return nil, err @@ -29,12 +46,19 @@ func New(port, target string) (*Client, error) { return nil, err } - return &Client{ + c := &Client{ laddr: laddr, raddr: raddr, mutex: sync.RWMutex{}, sessions: map[string]*session{}, - }, nil + timeout: 5 * time.Minute, + } + + for _, o := range options { + o(c) + } + + return c, nil } func (c *Client) ListenAndServe() error { @@ -76,8 +100,9 @@ func (c *Client) pruneSessions() { for range ticker.C { for _, session := range c.sessions { c.mutex.RLock() - if time.Since(session.updateTime) > time.Minute*5 { + if time.Since(session.updateTime) > c.timeout { delete(c.sessions, session.caddr.String()) + log.Tracef("session for %s deleted", session.caddr) } c.mutex.RUnlock() }