mk6i-open-oscar-server/server/toc/server.go

package toc

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/http"
	"sync"
	"syscall"
	"time"

	"github.com/patrickmn/go-cache"
	"golang.org/x/sync/errgroup"
	"golang.org/x/time/rate"

	"github.com/mk6i/open-oscar-server/wire"
)

var (
	// errClientReq indicates that an error occurred while reading a client request
	errClientReq = errors.New("failed to read client request")

	// errServerWrite indicates that an error occurred while writing a server response
	errServerWrite = errors.New("failed to send server response")

	// errTOCProcessing indicates that an error occurred in the TOC handler
	errTOCProcessing = errors.New("failed to process TOC request")
)

type contextKey string

const (
	screenNameContextKey contextKey = "screenName"
)

// bufferedConn is a wrapper around net.Conn that allows peeking into the
// incoming connection without consuming data. It is useful for multiplexing
// TOC/HTTP and TOC/FLAP connections.
//
// It embeds net.Conn, so all standard connection methods remain available.
type bufferedConn struct {
	r *bufio.Reader
	net.Conn
}

// newBufferedConn wraps a net.Conn with buffered reading capabilities.
func newBufferedConn(c net.Conn) bufferedConn {
	return bufferedConn{bufio.NewReader(c), c}
}

// Peek returns the next n bytes from the buffer without advancing the reader.
// If fewer than n bytes are available, it returns an error.
func (b bufferedConn) Peek(n int) ([]byte, error) {
	return b.r.Peek(n)
}

// Read reads data into p from the buffered connection.
// It prioritizes buffered data before reading from the underlying connection.
func (b bufferedConn) Read(p []byte) (int, error) {
	return b.r.Read(p)
}

// channelListener is an implementation of net.Listener that accepts connections
// from a channel instead of a network socket. It is useful for attaching an
// HTTP service to a connection on the fly.
type channelListener struct {
	ch  chan net.Conn // Channel used to receive connections.
	ctx context.Context
}

// Accept waits for and returns the next connection from the channel.
// If the channel is closed, it returns io.EOF to indicate no more connections.
func (l *channelListener) Accept() (net.Conn, error) {
	select {
	case <-l.ctx.Done():
		return nil, io.EOF
	case ch := <-l.ch:
		return ch, nil
	}
}

// Close closes the listener. Since channelListener does not manage an actual
// network connection, this is a no-op and always returns nil.
func (l *channelListener) Close() error {
	return nil
}

// Addr returns the network address of the listener.
// Since channelListener is not bound to a real network address, it returns nil.
func (l *channelListener) Addr() net.Addr {
	return nil
}

// IPRateLimiter provides per-IP rate limiting using a token bucket algorithm.
// It caches individual rate limiters per IP address with automatic TTL expiration.
type IPRateLimiter struct {
	cache *cache.Cache // In-memory cache of rate limiters keyed by IP
	rate  rate.Limit   // Allowed request rate (events per second)
	burst int          // Maximum burst size
}

// NewIPRateLimiter returns a new IPRateLimiter that limits each IP to the specified
// rate and burst, with limiter state expiring after the given TTL.
// Entries are retained for up to 2×TTL to reduce churn under frequent lookups.
func NewIPRateLimiter(rate rate.Limit, burst int, ttl time.Duration) *IPRateLimiter {
	return &IPRateLimiter{
		cache: cache.New(ttl, 2*ttl),
		rate:  rate,
		burst: burst,
	}
}

// Allow returns true if the request from the given IP is allowed under its rate limit.
// If no limiter exists for the IP, one is created and tracked in the cache.
func (l *IPRateLimiter) Allow(ip string) (allowed bool) {
	limiter, found := l.cache.Get(ip)
	if !found {
		limiter = rate.NewLimiter(l.rate, l.burst)
		l.cache.Set(ip, limiter, cache.DefaultExpiration)
	}
	return limiter.(*rate.Limiter).Allow()
}

func NewServer(
	endpointCfg []string,
	logger *slog.Logger,
	BOSProxy OSCARProxy,
) *Server {

	ctx, cancel := context.WithCancel(context.Background())

	s := &Server{
		bosProxy:       BOSProxy,
		conns:          make(map[net.Conn]struct{}),
		endpointCfg:    endpointCfg,
		logger:         logger,
		servers:        make([]*http.Server, 0, len(endpointCfg)),
		shutdownCancel: cancel,
		shutdownCtx:    ctx,
	}

	for range endpointCfg {
		s.servers = append(s.servers, &http.Server{
			Handler: BOSProxy.NewServeMux(),
			BaseContext: func(net.Listener) context.Context {
				return s.shutdownCtx
			},
		})
	}

	return s
}

// Server implements a TOC protocol server that multiplexes TOC/HTTP and
// TOC/FLAP requests. It acts as a gateway, forwarding all TOC requests
// to the OSCAR server for processing.
type Server struct {
	bosProxy OSCARProxy
	logger   *slog.Logger

	endpointCfg []string
	listeners   []net.Listener
	servers     []*http.Server

	connMu sync.Mutex
	conns  map[net.Conn]struct{}

	connWg   sync.WaitGroup
	listenWg sync.WaitGroup

	shutdownCtx    context.Context
	shutdownCancel context.CancelFunc
}

func (s *Server) ListenAndServe() error {
	g, ctx := errgroup.WithContext(s.shutdownCtx)

	for i, cfg := range s.endpointCfg {
		ln, err := net.Listen("tcp", cfg)
		if err != nil {
			s.cleanupListeners()
			s.shutdownCancel()
			return fmt.Errorf("unable to start TOC server: %w", err)
		}

		s.logger.InfoContext(ctx, "starting server", "listen_host", cfg)

		s.listeners = append(s.listeners, ln)
		s.listenWg.Add(1)

		httpCh := make(chan net.Conn)

		g.Go(func() error {
			cl := &channelListener{
				ch:  httpCh,
				ctx: s.shutdownCtx,
			}
			if err := s.servers[i].Serve(cl); !errors.Is(err, http.ErrServerClosed) && !errors.Is(err, io.EOF) {
				s.shutdownCancel()
				return err
			}
			return nil
		})

		g.Go(func() error {
			s.acceptLoop(ctx, ln, httpCh)
			return nil
		})
	}

	return g.Wait()
}

func (s *Server) Shutdown(ctx context.Context) error {
	s.logger.Debug("Initiating graceful shutdown...")
	s.shutdownCancel()
	s.cleanupListeners()

	// Wait for handlers to complete
	done := make(chan struct{})
	go func() {
		s.connWg.Wait()
		s.listenWg.Wait()
		close(done)
	}()

	for _, srv := range s.servers {
		_ = srv.Shutdown(ctx)
	}

	select {
	case <-done:
		s.logger.Info("shutdown complete")
	case <-ctx.Done():
		s.logger.Info("shutdown complete, but connections didn't close cleanly")
	}

	return nil
}

func (s *Server) cleanupListeners() {
	for _, ln := range s.listeners {
		_ = ln.Close()
	}
	s.listeners = nil
}

func (s *Server) acceptLoop(ctx context.Context, ln net.Listener, httpCh chan net.Conn) {
	defer s.listenWg.Done()

	for {
		conn, err := ln.Accept()
		if err != nil {
			if errors.Is(err, net.ErrClosed) {
				return
			}
			s.logger.Error("accept error", "err", err.Error())
			continue
		}

		go func() {
			if err := s.handleConnection(conn, ctx, httpCh); err != nil {
				s.logger.InfoContext(ctx, "user session failed", "err", err.Error())
			}
		}()
	}
}

// handleConnection inspects and routes an incoming connection. If the connection
// starts with "FLAP", handle as TOC/FLAP; otherwise, dispatch for HTTP
// processing.
func (s *Server) handleConnection(conn net.Conn, ctx context.Context, httpCh chan net.Conn) error {
	bufCon := newBufferedConn(conn)

	doFlap := "FLAP"
	buf, err := bufCon.Peek(len(doFlap))
	if err != nil {
		return fmt.Errorf("bufCon.Peek: %w", err)
	}

	// handle TOC/FLAP
	if string(buf) == doFlap {
		defer func() {
			// untrack connections
			s.connMu.Lock()
			delete(s.conns, conn)
			s.connMu.Unlock()

			_ = conn.Close()
			s.connWg.Done()
		}()

		// track connection
		s.connMu.Lock()
		s.conns[conn] = struct{}{}
		s.connMu.Unlock()

		s.connWg.Add(1)

		if err = s.dispatchFLAP(ctx, bufCon); err != nil {
			switch {
			case errors.Is(err, io.EOF):
			case errors.Is(err, net.ErrClosed):
			case errors.Is(err, syscall.ECONNRESET):
				return nil
			default:
				return fmt.Errorf("s.dispatchFLAP: %w", err)
			}
		}
		return nil
	}

	// handle TOC/HTTP
	select {
	case httpCh <- bufCon:
		return nil
	case <-ctx.Done():
		return nil
	}
}

func (s *Server) dispatchFLAP(ctx context.Context, conn net.Conn) error {
	var once sync.Once

	closeConn := func() {
		once.Do(func() {
			_ = conn.Close()
		})
	}
	defer closeConn()

	clientFlap, err := s.initFLAP(ctx, conn)
	if err != nil {
		return err
	}

	me, err := s.login(ctx, clientFlap, conn.RemoteAddr())
	if err != nil {
		return fmt.Errorf("s.login: %w", err)
	}
	if me == nil {
		return nil // signon failed
	}

	ctx = context.WithValue(ctx, screenNameContextKey, me.IdentScreenName())

	return s.handleTOCRequest(ctx, closeConn, me, clientFlap)
}

// handleTOCRequest processes incoming TOC requests and coordinates their handling.
// It reads client requests, processes TOC commands, and sends responses back to the client.
//
// Returns:
//   - errClientReq if an error occurs while reading the TOC request. wraps
//     io.EOF if the client disconnected.
//   - errTOCProcessing if an error occurs while processing the TOC command.
//   - errServerWrite if an error occurs while sending the TOC response.
func (s *Server) handleTOCRequest(
	ctx context.Context,
	closeConn func(),
	me *tocConn,
	clientFlap *wire.FlapClient,
) error {
	// TOC response queue
	msgCh := make(chan []string, 1)

	g, ctx := errgroup.WithContext(ctx)

	me.toCh = msgCh
	me.goAsync = g.Go

	// process TOC client requests and enqueue TOC server responses
	g.Go(func() error {
		err := s.runClientCommands(ctx, me, clientFlap)
		return errors.Join(err, errClientReq)
	})

	// translate OSCAR server responses to TOC responses and enqueue them
	g.Go(func() error {
		err := s.bosProxy.RecvBOS(ctx, me)
		closeConn() // unblock runClientCommands
		return errors.Join(err, errTOCProcessing)
	})

	// send TOC server responses to the client
	g.Go(func() error {
		err := s.sendToClient(ctx, msgCh, clientFlap)
		closeConn() // unblock runClientCommands
		return errors.Join(err, errServerWrite)
	})

	return g.Wait()
}

func (s *Server) runClientCommands(ctx context.Context, me *tocConn, clientFlap *wire.FlapClient) error {
	for {
		clientFrame, err := clientFlap.ReceiveFLAP()
		if err != nil {
			return err
		}
		switch clientFrame.FrameType {
		case wire.FLAPFrameSignoff:
			return io.EOF // client disconnected
		case wire.FLAPFrameKeepAlive:
			// keep alive heartbeat, do nothing for now.
			// todo set connection deadline to future time
		case wire.FLAPFrameData:
			clientFrame.Payload = bytes.TrimRight(clientFrame.Payload, "\x00") // trim null terminator

			if len(clientFrame.Payload) == 0 {
				return errors.New("TOC command is empty")
			}
			if len(clientFrame.Payload) > 2048 {
				return errors.New("TOC command exceeds maximum length (2048)")
			}

			msg := s.bosProxy.RecvClientCmd(ctx, me, clientFrame.Payload)
			if len(msg) > 0 {
				select {
				case me.toCh <- msg:
				case <-ctx.Done():
					return nil
				}
			}
		default:
			return fmt.Errorf("unexpected clientFlap clientFrame type %d", clientFrame.FrameType)
		}
	}
}

func (s *Server) sendToClient(ctx context.Context, toClient <-chan []string, clientFlap *wire.FlapClient) error {
	for {
		select {
		case <-ctx.Done():
			return nil
		case msgs := <-toClient:
			for _, m := range msgs {
				if err := clientFlap.SendDataFrame([]byte(m)); err != nil {
					return fmt.Errorf("clientFlap.SendDataFrame: %w", err)
				}
				if s.logger.Enabled(ctx, slog.LevelDebug) {
					s.logger.DebugContext(ctx, "server response", "command", m)
				} else {
					// just log the command, omit params
					idx := len(m)
					if col := bytes.IndexByte([]byte(m), ':'); col > -1 {
						idx = col
					}
					s.logger.InfoContext(ctx, "server response", "command", m[0:idx])
				}
			}
		}
	}
}

func (s *Server) login(ctx context.Context, clientFlap *wire.FlapClient, clientAddr net.Addr) (*tocConn, error) {
	clientFrame, err := clientFlap.ReceiveFLAP()
	if err != nil {
		if errors.Is(err, io.EOF) {
			return nil, nil
		}
		return nil, fmt.Errorf("clientFlap.ReceiveFLAP: %w", err)
	}

	me, reply := s.bosProxy.Signon(ctx, clientFrame.Payload, clientAddr)

	for _, m := range reply {
		if err := clientFlap.SendDataFrame([]byte(m)); err != nil {
			if me != nil {
				me.bos.Close()
			}
			return nil, fmt.Errorf("clientFlap.SendDataFrame: %w", err)
		}
	}

	return me, nil
}

// initFLAP sets up a new FLAP connection. It returns a flap client if the
// connection successfully initialized. It accepts either "FLAPON\n\n" or
// "FLAPON\r\n\r\n".
func (s *Server) initFLAP(ctx context.Context, rw io.ReadWriter) (*wire.FlapClient, error) {
	buf := make([]byte, 8)

	_, err := io.ReadFull(rw, buf)
	if err != nil {
		return nil, fmt.Errorf("io.ReadFull: %w", err)
	}
	if string(buf[:6]) != "FLAPON" {
		return nil, fmt.Errorf("expected FLAPON, got %s", buf)
	}
	if buf[6] == '\r' && buf[7] == '\n' {
		crlf := make([]byte, 2)
		_, err := io.ReadFull(rw, crlf)
		if err != nil {
			return nil, fmt.Errorf("io.ReadFull: %w", err)
		}
		if crlf[0] != '\r' || crlf[1] != '\n' {
			return nil, fmt.Errorf("expected \\r\\n after FLAPON\\r\\n, got %s", crlf)
		}
	} else if buf[6] != '\n' || buf[7] != '\n' {
		return nil, fmt.Errorf("expected FLAPON then \\n\\n or \\r\\n\\r\\n, got %s", buf)
	}

	clientFlap := wire.NewFlapClient(0, rw, rw)

	if err := clientFlap.SendSignonFrame(nil); err != nil {
		return nil, fmt.Errorf("clientFlap.SendSignonFrame: %w", err)
	}

	frame, err := clientFlap.ReceiveSignonFrame()
	if err != nil {
		return nil, fmt.Errorf("clientFlap.ReceiveSignonFrame: %w", err)
	}
	if sn, hasSn := frame.String(0x01); hasSn {
		s.logger.DebugContext(ctx, "new connection", "screen_name", sn)
	} else {
		s.logger.DebugContext(ctx, "new connection from unknown screen name")
	}

	return clientFlap, nil
}