server.go

// SPDX-License-Identifier: Apache-2.0

package frisbee

import (
	"context"
	"crypto/tls"
	"errors"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/loopholelabs/frisbee-go/pkg/packet"
	"github.com/loopholelabs/logging/types"
)

var (
	OnClosedNil = errors.New("OnClosed function cannot be nil")
	PreWriteNil = errors.New("PreWrite function cannot be nil")
	ListenerNil = errors.New("listener cannot be nil")
)

var (
	defaultOnClosed = func(_ *Async, _ error) {}

	defaultPreWrite = func() {}

	defaultStreamHandler = func(stream *Stream) {
		_ = stream.Close()
	}
)

// Server accepts connections from frisbee Clients and can send and receive frisbee Packets
type Server struct {
	listener      net.Listener
	handlerTable  HandlerTable
	shutdown      atomic.Bool
	options       *Options
	wg            sync.WaitGroup
	connections   map[*Async]struct{}
	connectionsMu sync.Mutex
	startedCh     chan struct{}
	concurrency   uint64
	limiter       chan struct{}

	baseContext       context.Context
	baseContextCancel context.CancelFunc

	// onClosed is a function run by the server whenever a connection is closed
	onClosed func(*Async, error)

	// preWrite is run by the server before a write happens
	preWrite func()

	// streamHandler is used to handle incoming client-initiated streams on the server
	streamHandler func(*Stream)

	// ConnContext is used to define a connection-specific context based on the incoming connection
	// and is run whenever a new connection is opened
	ConnContext func(context.Context, *Async) context.Context

	// StreamContext is used to define a stream-specific context based on the incoming stream
	// and is run whenever a new stream is opened
	StreamContext func(context.Context, *Stream) context.Context

	// PacketContext is used to define a handler-specific contexts based on the incoming packet
	// and is run whenever a new packet arrives
	PacketContext func(context.Context, *packet.Packet) context.Context

	// UpdateContext is used to update a handler-specific context whenever the returned
	// Action from a handler is UPDATE
	UpdateContext func(context.Context, *Async) context.Context
}

// NewServer returns an uninitialized frisbee Server with the registered HandlerTable.
// The Start method must then be called to start the server and listen for connections.
func NewServer(handlerTable HandlerTable, ctx context.Context, opts ...Option) (*Server, error) {
	options := loadOptions(opts...)

	baseContext, baseContextCancel := context.WithCancel(ctx)

	s := &Server{
		options:           options,
		connections:       make(map[*Async]struct{}),
		startedCh:         make(chan struct{}),
		baseContext:       baseContext,
		baseContextCancel: baseContextCancel,
		onClosed:          defaultOnClosed,
		preWrite:          defaultPreWrite,
		streamHandler:     defaultStreamHandler,
	}

	return s, s.SetHandlerTable(handlerTable)
}

// SetOnClosed sets the onClosed function for the server. If f is nil, it returns an error.
func (s *Server) SetOnClosed(f func(*Async, error)) error {
	if f == nil {
		return OnClosedNil
	}
	s.onClosed = f
	return nil
}

// SetPreWrite sets the preWrite function for the server. If f is nil, it returns an error.
func (s *Server) SetPreWrite(f func()) error {
	if f == nil {
		return PreWriteNil
	}
	s.preWrite = f
	return nil
}

// SetStreamHandler sets the streamHandler function for the server.
func (s *Server) SetStreamHandler(f func(context.Context, *Stream)) error {
	s.streamHandler = func(stream *Stream) {
		streamCtx := s.baseContext
		if s.StreamContext != nil {
			streamCtx = s.StreamContext(streamCtx, stream)
		}
		f(streamCtx, stream)
	}
	return nil
}

// SetHandlerTable sets the handler table for the server.
//
// This function should not be called once the server has started.
func (s *Server) SetHandlerTable(handlerTable HandlerTable) error {
	for i := uint16(0); i < RESERVED9; i++ {
		if _, ok := handlerTable[i]; ok {
			return InvalidHandlerTable
		}
	}

	s.handlerTable = handlerTable
	return nil
}

// GetHandlerTable gets the handler table for the server.
//
// This function should not be called once the server has started.
func (s *Server) GetHandlerTable() HandlerTable {
	return s.handlerTable
}

// SetConcurrency sets the maximum number of concurrent goroutines that will be created
// by the server to handle incoming packets.
//
// An important caveat of this is that handlers must always thread-safe if they share resources
// between connections. If the concurrency is set to a value != 1, then the handlers
// must also be thread-safe if they share resources per connection.
//
// This function should not be called once the server has started.
func (s *Server) SetConcurrency(concurrency uint64) {
	s.concurrency = concurrency
	if s.concurrency > 1 {
		s.limiter = make(chan struct{}, s.concurrency)
	}
}

// Start will start the frisbee server and its reactor goroutines
// to receive and handle incoming connections. If the baseContext, ConnContext,
// onClosed, OnShutdown, or preWrite functions have not been defined, it will
// use the default functions for these.
func (s *Server) Start(addr string) error {
	var listener net.Listener
	var err error
	if s.options.TLSConfig != nil {
		listener, err = tls.Listen("tcp", addr, s.options.TLSConfig)
	} else {
		listener, err = net.Listen("tcp", addr)
	}
	if err != nil {
		return err
	}
	return s.StartWithListener(listener)
}

// StartWithListener will start the frisbee server and its reactor goroutines
// to receive and handle incoming connections with a given net.Listener. If the baseContext, ConnContext,
// onClosed, OnShutdown, or preWrite functions have not been defined, it will
// use the default functions for these.
func (s *Server) StartWithListener(listener net.Listener) error {
	if listener == nil {
		return ListenerNil
	}
	s.listener = listener
	s.wg.Add(1)
	close(s.startedCh)
	return s.handleListener()
}

// started returns a channel that will be closed when the server has successfully started
//
// This is meant to only be used for testing purposes.
func (s *Server) started() <-chan struct{} {
	return s.startedCh
}

func (s *Server) handleListener() error {
	var backoff time.Duration
	for {
		newConn, err := s.listener.Accept()
		if err != nil {
			if s.shutdown.Load() {
				s.wg.Done()
				return nil
			}
			if ne, ok := err.(temporary); ok && ne.Temporary() {
				if backoff == 0 {
					backoff = minBackoff
				} else {
					backoff *= 2
				}
				if backoff > maxBackoff {
					backoff = maxBackoff
				}
				s.Logger().Warn().Err(err).Msgf("Temporary Accept Error, retrying in %s", backoff)
				time.Sleep(backoff)
				if s.shutdown.Load() {
					s.wg.Done()
					return nil
				}
				continue
			}
			s.wg.Done()
			return err
		}
		backoff = 0

		s.ServeConn(newConn)
	}
}

func (s *Server) createHandler(conn *Async, closed *atomic.Bool, wg *sync.WaitGroup, ctx context.Context, cancel context.CancelFunc) func(*packet.Packet) {
	return func(p *packet.Packet) {
		handlerFunc := s.handlerTable[p.Metadata.Operation]
		if handlerFunc != nil {
			packetCtx := ctx
			if s.PacketContext != nil {
				packetCtx = s.PacketContext(packetCtx, p)
			}
			outgoing, action := handlerFunc(packetCtx, p)
			if outgoing != nil && outgoing.Metadata.ContentLength == uint32(outgoing.Content.Len()) {
				s.preWrite()
				err := conn.WritePacket(outgoing)
				if outgoing != p {
					packet.Put(outgoing)
				}
				packet.Put(p)
				if err != nil {
					_ = conn.Close()
					if closed.CompareAndSwap(false, true) {
						s.onClosed(conn, err)
					}
					cancel()
					wg.Done()
					return
				}
			} else {
				packet.Put(p)
			}
			switch action {
			case NONE:
			case CLOSE:
				_ = conn.Close()
				if closed.CompareAndSwap(false, true) {
					s.onClosed(conn, nil)
				}
				cancel()
			}
		} else {
			packet.Put(p)
		}
		wg.Done()
	}
}

func (s *Server) handleSinglePacket(frisbeeConn *Async, connCtx context.Context) {
	var p *packet.Packet
	var outgoing *packet.Packet
	var action Action
	var handlerFunc Handler
	var err error
	p, err = frisbeeConn.ReadPacket()
	if err != nil {
		_ = frisbeeConn.Close()
		s.onClosed(frisbeeConn, err)
		return
	}
	for {
		handlerFunc = s.handlerTable[p.Metadata.Operation]
		if handlerFunc != nil {
			packetCtx := connCtx
			if s.PacketContext != nil {
				packetCtx = s.PacketContext(packetCtx, p)
			}
			outgoing, action = handlerFunc(packetCtx, p)
			if outgoing != nil && outgoing.Metadata.ContentLength == uint32(outgoing.Content.Len()) {
				s.preWrite()
				err = frisbeeConn.WritePacket(outgoing)
				if outgoing != p {
					packet.Put(outgoing)
				}
				packet.Put(p)
				if err != nil {
					_ = frisbeeConn.Close()
					s.onClosed(frisbeeConn, err)
					return
				}
			} else {
				packet.Put(p)
			}
			switch action {
			case NONE:
			case CLOSE:
				_ = frisbeeConn.Close()
				s.onClosed(frisbeeConn, nil)
				return
			}
		} else {
			packet.Put(p)
		}
		p, err = frisbeeConn.ReadPacket()
		if err != nil {
			_ = frisbeeConn.Close()
			s.onClosed(frisbeeConn, err)
			return
		}
	}
}

func (s *Server) handleUnlimitedPacket(frisbeeConn *Async, connCtx context.Context) {
	p, err := frisbeeConn.ReadPacket()
	if err != nil {
		_ = frisbeeConn.Close()
		s.onClosed(frisbeeConn, err)
		return
	}
	wg := new(sync.WaitGroup)
	var closed atomic.Bool
	connCtx, cancel := context.WithCancel(connCtx)
	handle := s.createHandler(frisbeeConn, &closed, wg, connCtx, cancel)
	for {
		wg.Add(1)
		go handle(p)
		p, err = frisbeeConn.ReadPacket()
		if err != nil {
			_ = frisbeeConn.Close()
			if closed.CompareAndSwap(false, true) {
				s.onClosed(frisbeeConn, err)
			}
			cancel()
			wg.Wait()
			return
		}
	}
}

func (s *Server) handleLimitedPacket(frisbeeConn *Async, connCtx context.Context) {
	p, err := frisbeeConn.ReadPacket()
	if err != nil {
		_ = frisbeeConn.Close()
		s.onClosed(frisbeeConn, err)
		return
	}
	wg := new(sync.WaitGroup)
	var closed atomic.Bool
	connCtx, cancel := context.WithCancel(connCtx)
	handler := s.createHandler(frisbeeConn, &closed, wg, connCtx, cancel)
	handle := func(p *packet.Packet) {
		handler(p)
		<-s.limiter
	}
	for {
		select {
		case s.limiter <- struct{}{}:
			wg.Add(1)
			go handle(p)
			p, err = frisbeeConn.ReadPacket()
			if err != nil {
				_ = frisbeeConn.Close()
				if closed.CompareAndSwap(false, true) {
					s.onClosed(frisbeeConn, err)
				}
				cancel()
				wg.Wait()
				return
			}
		case <-connCtx.Done():
			_ = frisbeeConn.Close()
			if closed.CompareAndSwap(false, true) {
				s.onClosed(frisbeeConn, err)
			}
			wg.Wait()
			return
		}
	}
}

// ServeConn takes a net.Conn and starts a goroutine to handle it using the Server.
func (s *Server) ServeConn(conn net.Conn) {
	s.wg.Add(1)
	go s.serveConn(conn)
}

// serveConn takes a net.Conn and serves it using the Server
// and assumes that the server's wait group has been incremented by 1.
func (s *Server) serveConn(newConn net.Conn) {
	var err error
	switch v := newConn.(type) {
	case *net.TCPConn:
		err = v.SetKeepAlive(true)
		if err != nil {
			s.Logger().Error().Err(err).Msg("Error while setting TCP Keepalive")
			_ = v.Close()
			s.wg.Done()
			return
		}
		err = v.SetKeepAlivePeriod(s.options.KeepAlive)
		if err != nil {
			s.Logger().Error().Err(err).Msg("Error while setting TCP Keepalive Period")
			_ = v.Close()
			s.wg.Done()
			return
		}
	}

	frisbeeConn := NewAsync(newConn, s.Logger(), s.streamHandler)
	connCtx := s.baseContext
	s.connectionsMu.Lock()
	if s.shutdown.Load() {
		s.wg.Done()
		return
	}
	s.connections[frisbeeConn] = struct{}{}
	s.connectionsMu.Unlock()
	if s.ConnContext != nil {
		connCtx = s.ConnContext(connCtx, frisbeeConn)
	}
	switch s.concurrency {
	case 0:
		s.handleUnlimitedPacket(frisbeeConn, connCtx)
	case 1:
		s.handleSinglePacket(frisbeeConn, connCtx)
	default:
		s.handleLimitedPacket(frisbeeConn, connCtx)
	}
	s.connectionsMu.Lock()
	if !s.shutdown.Load() {
		delete(s.connections, frisbeeConn)
	}
	s.connectionsMu.Unlock()
	s.wg.Done()
}

// Logger returns the server's logger (useful for ServerRouter functions)
func (s *Server) Logger() types.Logger {
	return s.options.Logger
}

// Shutdown shuts down the frisbee server and kills all the goroutines and active connections
func (s *Server) Shutdown() error {
	if s.shutdown.CompareAndSwap(false, true) {
		s.baseContextCancel()
		s.connectionsMu.Lock()
		for c := range s.connections {
			_ = c.Close()
			delete(s.connections, c)
		}
		s.connectionsMu.Unlock()
		defer s.wg.Wait()
		if s.listener != nil {
			return s.listener.Close()
		}
	}
	return nil
}