rebound/server.go

package rebound

import (
	"log/slog"
	"net"
	"os"
	"time"

	"forge.cadoles.com/wpetit/rebound/http"
	"forge.cadoles.com/wpetit/rebound/ssh"
	"forge.cadoles.com/wpetit/rebound/stat"
	"github.com/pkg/errors"
)

type Server struct {
	listener net.Listener
	opts     *Options
	stats    *stat.Store
}

func (s *Server) Start() error {
	s.log("[INFO] listening on %s", s.opts.Address)

	if err := s.stats.Load(s.opts.StatsFile); err != nil {
		if errors.Is(err, os.ErrNotExist) {
			s.log("[INFO] stats file does not exist. ignoring.")
		} else {
			return errors.WithStack(err)
		}
	}

	listener, err := net.Listen("tcp", s.opts.Address)
	if err != nil {
		return errors.WithStack(err)
	}

	s.listener = listener

	sshListener, httpListener := s.muxListener(listener)

	go func() {
		defer listener.Close()

		server := ssh.NewServer(
			ssh.WithHostKey(s.opts.SSH.HostKey),
			ssh.WithPublicHost(s.opts.SSH.PublicHost),
			ssh.WithPublicPort(s.opts.SSH.PublicPort),
			ssh.WithSockDir(s.opts.SSH.SockDir),
			ssh.WithLogger(s.opts.SSH.Logger),
			ssh.WithStats(s.stats),
		)

		if err := server.Serve(sshListener); err != nil {
			s.log("[ERROR] %+v", errors.WithStack(err))
			listener.Close()
		}
	}()

	go func() {
		defer listener.Close()

		server := http.NewServer(
			http.WithCustomDir(s.opts.HTTP.CustomDir),
			http.WithTemplateData(s.opts.HTTP.TemplateData),
			http.WithLogger(s.opts.HTTP.Logger),
			http.WithStats(s.stats),
		)

		if err := server.Serve(httpListener); err != nil {
			s.log("[ERROR] %+v", errors.WithStack(err))

		}
	}()

	go func() {
		defer listener.Close()

		ticker := time.NewTicker(s.opts.StatsFileSaveInterval)

		for {
			<-ticker.C

			slog.Info("saving stats", slog.String("file", s.opts.StatsFile), slog.Duration("interval", s.opts.StatsFileSaveInterval))
			if err := s.stats.Save(s.opts.StatsFile); err != nil {
				slog.Error("could not save stat file", slog.Any("error", errors.WithStack(err)))

				return
			}
		}
	}()

	return nil
}

func (s *Server) Stop() error {
	if s.listener == nil {
		return nil
	}

	if err := s.listener.Close(); err != nil {
		return errors.WithStack(err)
	}

	s.listener = nil

	return nil
}

func (s *Server) log(message string, args ...any) {
	s.opts.Logger(message, args...)
}

func NewServer(funcs ...OptionFunc) *Server {
	opts := DefaultOptions()
	for _, fn := range funcs {
		fn(opts)
	}

	return &Server{
		opts:  opts,
		stats: stat.NewStore(),
	}
}