Moonshark/moonshark.go

package main

import (
	"Moonshark/config"
	"Moonshark/http"
	"Moonshark/logger"
	"Moonshark/metadata"
	"Moonshark/router"
	"Moonshark/runner"
	"Moonshark/sessions"
	"Moonshark/utils"
	"Moonshark/watchers"
	"bytes"
	"context"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"

	color "git.sharkk.net/Go/Color"

	fin "git.sharkk.net/Sharkk/Fin"
	"github.com/valyala/fasthttp"
)

var (
	cfg    *config.Config           // Server config from Fin file
	rtr    *router.Router           // Lua file router
	rnr    *runner.Runner           // Lua runner
	svr    *fasthttp.Server         // FastHTTP server
	pub    fasthttp.RequestHandler  // Public asset handler
	snm    *sessions.SessionManager // Session data manager
	wmg    *watchers.WatcherManager // Watcher manager
	dbg    bool                     // Debug mode flag
	pubPfx []byte                   // Cached public asset prefix
)

func main() {
	cfgPath := flag.String("config", "config", "Path to Fin config file")
	dbgFlag := flag.Bool("debug", false, "Force debug mode")
	sptPath := flag.String("script", "", "Path to Lua script to execute once")
	flag.Parse()

	// Init sequence
	sptMode := *sptPath != ""
	color.SetColors(color.DetectShellColors())
	banner(sptMode)

	// Load Fin-based config
	cfg = config.New(fin.LoadFromFile(*cfgPath))

	// Setup debug mode
	dbg = *dbgFlag || cfg.Server.Debug
	logger.Debug(dbg)
	logger.Debugf("Debug logging enabled") // Only prints if dbg is true
	utils.Debug(dbg)                       // @TODO find a better way to do this

	// Determine Lua runner pool size
	poolSize := cfg.Runner.PoolSize
	if sptMode {
		poolSize = 1
	}

	// Set up the Lua runner
	if err := initRunner(poolSize); err != nil {
		logger.Fatalf("Runner failed to init: %v", err)
	}

	// If in script mode, attempt to run the Lua script at the given path
	if sptMode {
		if err := handleScriptMode(*sptPath); err != nil {
			logger.Fatalf("Script execution failed: %v", err)
		}

		shutdown()
		return
	}

	// Set up the Lua router
	if err := initRouter(); err != nil {
		logger.Fatalf("Router failed to init: %s", color.Red(err.Error()))
	}

	// Set up the file watcher manager
	if err := setupWatchers(); err != nil {
		logger.Fatalf("Watcher manager failed to init: %s", color.Red(err.Error()))
	}

	// Set up the HTTP portion of the server
	logger.Http(cfg.Server.HTTPLogging) // Whether we'll log HTTP request results
	svr = http.NewHttpServer(cfg, requestMux, dbg)
	pub = http.NewPublicHandler(cfg.Dirs.Public, cfg.Server.PublicPrefix)
	pubPfx = []byte(cfg.Server.PublicPrefix) // Avoids casting to []byte when check prefixes
	snm = sessions.NewSessionManager(sessions.DefaultMaxSessions)

	// Start the HTTP server
	logger.Infof("Surf's up on port %s!", color.Cyan(strconv.Itoa(cfg.Server.Port)))
	go func() {
		if err := svr.ListenAndServe(":" + strconv.Itoa(cfg.Server.Port)); err != nil {
			if err.Error() != "http: Server closed" {
				logger.Errorf("Server error: %v", err)
			}
		}
	}()

	// Handle a shutdown signal
	stop := make(chan os.Signal, 1)
	signal.Notify(stop, os.Interrupt, syscall.SIGTERM)
	<-stop

	fmt.Print("\n")
	logger.Infof("Shutdown signal received")
	shutdown()
}

// This is the primary request handler mux - determines whether we need to handle a Lua
// route or if we're serving a static file.
func requestMux(ctx *fasthttp.RequestCtx) {
	start := time.Now()
	method := ctx.Method()
	path := ctx.Path()

	// Handle static file request
	if bytes.HasPrefix(path, pubPfx) {
		pub(ctx)
		logRequest(ctx, method, path, start)
		return
	}

	// See if the requested route even exists
	bytecode, params, found := rtr.Lookup(string(method), string(path))
	if !found {
		http.Send404(ctx)
		logRequest(ctx, method, path, start)
		return
	}

	// If there's no bytecode then it's an internal server error
	if len(bytecode) == 0 {
		http.Send500(ctx, nil)
		logRequest(ctx, method, path, start)
	}

	// We've made it this far so the endpoint will likely load. Let's get any session data
	// for this request
	session := snm.GetSessionFromRequest(ctx)

	// Let's build an HTTP context for the Lua runner to consume
	luaCtx := runner.NewHTTPContext(ctx, params, session)
	defer luaCtx.Release()

	// Ask the runner to execute our endpoint with our context
	res, err := rnr.Execute(bytecode, luaCtx)
	if err != nil {
		logger.Errorf("Lua execution error: %v", err)
		http.Send500(ctx, err)
		logRequest(ctx, method, path, start)
		return
	}

	// Sweet, our execution went through! Let's now use the Response we got and build the HTTP response, then return
	// the response object to be cleaned. After, we'll log our request cus we are *done*
	applyResponse(ctx, res, session)
	runner.ReleaseResponse(res)
	logRequest(ctx, method, path, start)
}

func applyResponse(ctx *fasthttp.RequestCtx, resp *runner.Response, session *sessions.Session) {
	// Handle session updates
	if len(resp.SessionData) > 0 {
		if _, clearAll := resp.SessionData["__clear_all"]; clearAll {
			session.Clear()
			session.ClearFlash()
			delete(resp.SessionData, "__clear_all")
		}

		for k, v := range resp.SessionData {
			if v == "__DELETE__" {
				session.Delete(k)
			} else {
				session.Set(k, v)
			}
		}
	}

	// Handle flash data
	if flashData, ok := resp.Metadata["flash"].(map[string]any); ok {
		for k, v := range flashData {
			if err := session.FlashSafe(k, v); err != nil && dbg {
				logger.Warnf("Error setting flash data %s: %v", k, err)
			}
		}
	}

	// Apply session cookie
	snm.ApplySessionCookie(ctx, session)

	// Apply HTTP response
	http.ApplyResponse(resp, ctx)
}

// Attempts to start the Lua runner. poolSize allows overriding the config, like for script mode. A poolSize of
// 0 will default to the config, and if the config is 0 then it will default to GOMAXPROCS.
func initRunner(poolSize int) error {
	for _, dir := range cfg.Dirs.Libs {
		if !dirExists(dir) {
			logger.Warnf("Lib directory not found... %s", color.Yellow(dir))
		}
	}

	runner, err := runner.NewRunner(cfg, poolSize)
	if err != nil {
		return fmt.Errorf("lua runner init failed: %v", err)
	}
	rnr = runner

	logger.Infof("LuaRunner is g2g with %s states!", color.Yellow(strconv.Itoa(poolSize)))
	return nil
}

// Attempt to spin up the Lua router. Attempts to create the routes directory if it doesn't exist,
// since it's required for Moonshark to work.
func initRouter() error {
	if err := os.MkdirAll(cfg.Dirs.Routes, 0755); err != nil {
		return fmt.Errorf("failed to create routes directory: %w", err)
	}

	router, err := router.New(cfg.Dirs.Routes)
	if err != nil {
		return fmt.Errorf("lua router init failed: %v", err)
	}
	rtr = router

	logger.Infof("LuaRouter is g2g! %s", color.Yellow(cfg.Dirs.Routes))
	return nil
}

// Set up the file watchers.
func setupWatchers() error {
	wmg = watchers.NewWatcherManager()

	// Router watcher
	err := wmg.WatchDirectory(watchers.WatcherConfig{
		Dir:       cfg.Dirs.Routes,
		Callback:  rtr.Refresh,
		Recursive: true,
	})
	if err != nil {
		return fmt.Errorf("failed to watch routes directory: %v", err)
	}

	logger.Infof("Started watching Lua routes! %s", color.Yellow(cfg.Dirs.Routes))

	// Libs watchers
	for _, dir := range cfg.Dirs.Libs {
		err := wmg.WatchDirectory(watchers.WatcherConfig{
			Dir: dir,
			Callback: func(changes []watchers.FileChange) error {
				for _, change := range changes {
					if !change.IsDeleted && strings.HasSuffix(change.Path, ".lua") {
						rnr.NotifyFileChanged(change.Path)
					}
				}
				return nil
			},
			Recursive: true,
		})
		if err != nil {
			return fmt.Errorf("failed to watch modules directory: %v", err)
		}

		logger.Infof("Started watching Lua modules! %s", color.Yellow(dir))
	}

	return nil
}

// Attempts to execute the Lua script at the given path inside a fully initialized sandbox environment. Handy
// for pre-launch tasks and the like.
func handleScriptMode(path string) error {
	path, err := filepath.Abs(path)
	if err != nil {
		return fmt.Errorf("failed to resolve script path: %v", err)
	}

	if _, err := os.Stat(path); os.IsNotExist(err) {
		return fmt.Errorf("script file not found: %s", path)
	}

	logger.Infof("Executing: %s", path)

	resp, err := rnr.RunScriptFile(path)
	if err != nil {
		return fmt.Errorf("execution failed: %v", err)
	}

	if resp != nil && resp.Body != nil {
		logger.Infof("Script result: %v", resp.Body)
	} else {
		logger.Infof("Script executed successfully (no return value)")
	}

	return nil
}

func shutdown() {
	logger.Infof("Shutting down...")

	// Close down the HTTP server
	if svr != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()

		if err := svr.ShutdownWithContext(ctx); err != nil {
			logger.Errorf("HTTP server shutdown error: %v", err)
		}
	}

	// Close down the Lua runner if it exists
	if rnr != nil {
		rnr.Close()
	}

	// Close down the watcher manager if it exists
	if wmg != nil {
		wmg.Close()
	}

	logger.Infof("Shutdown complete")
}

// Print our super-awesome banner with the current version!
func banner(scriptMode bool) {
	if scriptMode {
		fmt.Println(color.Blue(fmt.Sprintf("Moonshark %s << Script Mode >>", metadata.Version)))
		return
	}

	banner := `
   _____                         _________.__                  __
  /     \   ____   ____   ____  /   _____/|  |__ _____ _______|  | __
 /  \ /  \ /  _ \ /  _ \ /    \ \_____  \ |  |  \\__  \\_  __ \  |/ /
/    Y    (  <_> |  <_> )   |  \/        \|   Y  \/ __ \|  | \/    <
\____|__  /\____/ \____/|___|  /_______  /|___|  (____  /__|  |__|_ \ %s
        \/                   \/        \/      \/     \/           \/
	`
	fmt.Println(color.Blue(fmt.Sprintf(banner, metadata.Version)))
}

func dirExists(path string) bool {
	info, err := os.Stat(path)
	return err == nil && info.IsDir()
}

func logRequest(ctx *fasthttp.RequestCtx, method, path []byte, start time.Time) {
	logger.Request(ctx.Response.StatusCode(), method, path, time.Since(start))
}