Moonshark/http/server.go

package http

import (
	"bytes"
	"context"
	"strings"
	"sync"
	"time"

	"Moonshark/router"
	"Moonshark/runner"
	"Moonshark/runner/lualibs"
	"Moonshark/sessions"
	"Moonshark/utils"
	"Moonshark/utils/color"
	"Moonshark/utils/config"
	"Moonshark/utils/logger"
	"Moonshark/utils/metadata"

	"github.com/valyala/fasthttp"
)

var (
	//methodGET   = []byte("GET")
	methodPOST  = []byte("POST")
	methodPUT   = []byte("PUT")
	methodPATCH = []byte("PATCH")
	debugPath   = []byte("/debug/stats")
)

type Server struct {
	luaRouter         *router.LuaRouter
	staticHandler     fasthttp.RequestHandler
	staticFS          *fasthttp.FS
	luaRunner         *runner.Runner
	fasthttpServer    *fasthttp.Server
	debugMode         bool
	cfg               *config.Config
	sessionManager    *sessions.SessionManager
	errorConfig       utils.ErrorPageConfig
	ctxPool           sync.Pool
	paramsPool        sync.Pool
	staticPrefix      string
	staticPrefixBytes []byte

	// Cached error pages
	cached404    []byte
	cached500    []byte
	errorCacheMu sync.RWMutex
}

func New(luaRouter *router.LuaRouter, runner *runner.Runner, cfg *config.Config, debugMode bool) *Server {
	staticPrefix := cfg.Server.StaticPrefix
	if !strings.HasPrefix(staticPrefix, "/") {
		staticPrefix = "/" + staticPrefix
	}
	if !strings.HasSuffix(staticPrefix, "/") {
		staticPrefix = staticPrefix + "/"
	}

	s := &Server{
		luaRouter:         luaRouter,
		luaRunner:         runner,
		debugMode:         debugMode,
		cfg:               cfg,
		sessionManager:    sessions.GlobalSessionManager,
		staticPrefix:      staticPrefix,
		staticPrefixBytes: []byte(staticPrefix),
		errorConfig: utils.ErrorPageConfig{
			OverrideDir: cfg.Dirs.Override,
			DebugMode:   debugMode,
		},
		ctxPool: sync.Pool{
			New: func() any {
				return make(map[string]any, 6)
			},
		},
		paramsPool: sync.Pool{
			New: func() any {
				return make(map[string]any, 4)
			},
		},
	}

	// Pre-cache error pages
	s.cached404 = []byte(utils.NotFoundPage(s.errorConfig, ""))
	s.cached500 = []byte(utils.InternalErrorPage(s.errorConfig, "", "Internal Server Error"))

	// Setup static file serving
	if cfg.Dirs.Static != "" {
		s.staticFS = &fasthttp.FS{
			Root:                 cfg.Dirs.Static,
			IndexNames:           []string{"index.html"},
			GenerateIndexPages:   false,
			AcceptByteRange:      true,
			Compress:             true,
			CompressedFileSuffix: ".gz",
			CompressBrotli:       true,
			CompressZstd:         true,
			PathRewrite:          fasthttp.NewPathPrefixStripper(len(staticPrefix) - 1),
		}
		s.staticHandler = s.staticFS.NewRequestHandler()
	}

	s.fasthttpServer = &fasthttp.Server{
		Handler:                       s.handleRequest,
		Name:                          "Moonshark/" + metadata.Version,
		ReadTimeout:                   30 * time.Second,
		WriteTimeout:                  30 * time.Second,
		IdleTimeout:                   120 * time.Second,
		MaxRequestBodySize:            16 << 20,
		TCPKeepalive:                  true,
		TCPKeepalivePeriod:            60 * time.Second,
		ReduceMemoryUsage:             true,
		DisablePreParseMultipartForm:  true,
		DisableHeaderNamesNormalizing: true,
		NoDefaultServerHeader:         true,
		StreamRequestBody:             true,
	}

	return s
}

func (s *Server) ListenAndServe(addr string) error {
	logger.Infof("Catch the swell at %s", color.Cyan("http://localhost"+addr))
	return s.fasthttpServer.ListenAndServe(addr)
}

func (s *Server) Shutdown(ctx context.Context) error {
	return s.fasthttpServer.ShutdownWithContext(ctx)
}

func (s *Server) handleRequest(ctx *fasthttp.RequestCtx) {
	start := time.Now()
	methodBytes := ctx.Method()
	pathBytes := ctx.Path()

	if s.debugMode && bytes.Equal(pathBytes, debugPath) {
		s.handleDebugStats(ctx)
		if s.cfg.Server.HTTPLogging {
			logger.Request(ctx.Response.StatusCode(), string(methodBytes), string(pathBytes), time.Since(start))
		}
		return
	}

	if s.staticHandler != nil && bytes.HasPrefix(pathBytes, s.staticPrefixBytes) {
		s.staticHandler(ctx)
		if s.cfg.Server.HTTPLogging {
			logger.Request(ctx.Response.StatusCode(), string(methodBytes), string(pathBytes), time.Since(start))
		}
		return
	}

	bytecode, scriptPath, routeErr, params, found := s.luaRouter.GetRouteInfo(methodBytes, pathBytes)

	if found {
		if len(bytecode) == 0 || routeErr != nil {
			s.sendError(ctx, fasthttp.StatusInternalServerError, pathBytes, routeErr)
		} else {
			s.handleLuaRoute(ctx, bytecode, scriptPath, params, methodBytes, pathBytes)
		}
	} else {
		s.send404(ctx, pathBytes)
	}

	if s.cfg.Server.HTTPLogging {
		logger.Request(ctx.Response.StatusCode(), string(methodBytes), string(pathBytes), time.Since(start))
	}
}

func (s *Server) handleLuaRoute(ctx *fasthttp.RequestCtx, bytecode []byte, scriptPath string,
	params *router.Params, methodBytes, pathBytes []byte) {

	luaCtx := runner.NewHTTPContext(ctx)
	defer luaCtx.Release()

	if lualibs.GetGlobalEnvManager() != nil {
		luaCtx.Set("env", lualibs.GetGlobalEnvManager().GetAll())
	}

	sessionMap := s.ctxPool.Get().(map[string]any)
	defer func() {
		for k := range sessionMap {
			delete(sessionMap, k)
		}
		s.ctxPool.Put(sessionMap)
	}()

	session := s.sessionManager.GetSessionFromRequest(ctx)
	sessionMap["id"] = session.ID

	// Only get session data if not empty
	if !session.IsEmpty() {
		sessionMap["data"] = session.GetAll()
	} else {
		sessionMap["data"] = emptyMap
	}

	// Set basic context
	luaCtx.Set("method", string(methodBytes))
	luaCtx.Set("path", string(pathBytes))
	luaCtx.Set("host", string(ctx.Host()))
	luaCtx.Set("session", sessionMap)

	// Add headers to context
	headers := make(map[string]any)
	ctx.Request.Header.VisitAll(func(key, value []byte) {
		headers[string(key)] = string(value)
	})
	luaCtx.Set("headers", headers)

	// Handle params
	if params != nil && params.Count > 0 {
		paramMap := s.paramsPool.Get().(map[string]any)
		for i := range params.Count {
			paramMap[params.Keys[i]] = params.Values[i]
		}
		luaCtx.Set("params", paramMap)
		defer func() {
			for k := range paramMap {
				delete(paramMap, k)
			}
			s.paramsPool.Put(paramMap)
		}()
	} else {
		luaCtx.Set("params", emptyMap)
	}

	// Parse form data for POST/PUT/PATCH
	if bytes.Equal(methodBytes, methodPOST) ||
		bytes.Equal(methodBytes, methodPUT) ||
		bytes.Equal(methodBytes, methodPATCH) {
		if formData, err := ParseForm(ctx); err == nil {
			luaCtx.Set("form", formData)
		} else {
			if s.debugMode {
				logger.Warnf("Error parsing form: %v", err)
			}
			luaCtx.Set("form", emptyMap)
		}
	} else {
		luaCtx.Set("form", emptyMap)
	}

	response, err := s.luaRunner.Run(bytecode, luaCtx, scriptPath)
	if err != nil {
		logger.Errorf("Lua execution error: %v", err)
		s.sendError(ctx, fasthttp.StatusInternalServerError, pathBytes, err)
		return
	}

	// Handle session updates
	if len(response.SessionData) > 0 {
		if _, clearAll := response.SessionData["__clear_all"]; clearAll {
			session.Clear()
			delete(response.SessionData, "__clear_all")
		}

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

	s.sessionManager.ApplySessionCookie(ctx, session)
	runner.ApplyResponse(response, ctx)
	runner.ReleaseResponse(response)
}

func (s *Server) send404(ctx *fasthttp.RequestCtx, pathBytes []byte) {
	ctx.SetContentType("text/html; charset=utf-8")
	ctx.SetStatusCode(fasthttp.StatusNotFound)

	// Use cached 404 for common case
	if len(pathBytes) == 1 && pathBytes[0] == '/' {
		s.errorCacheMu.RLock()
		ctx.SetBody(s.cached404)
		s.errorCacheMu.RUnlock()
	} else {
		ctx.SetBody([]byte(utils.NotFoundPage(s.errorConfig, string(pathBytes))))
	}
}

func (s *Server) sendError(ctx *fasthttp.RequestCtx, status int, pathBytes []byte, err error) {
	ctx.SetContentType("text/html; charset=utf-8")
	ctx.SetStatusCode(status)

	if err == nil {
		s.errorCacheMu.RLock()
		ctx.SetBody(s.cached500)
		s.errorCacheMu.RUnlock()
	} else {
		ctx.SetBody([]byte(utils.InternalErrorPage(s.errorConfig, string(pathBytes), err.Error())))
	}
}

func (s *Server) handleDebugStats(ctx *fasthttp.RequestCtx) {
	stats := utils.CollectSystemStats(s.cfg)
	routeCount, bytecodeBytes := s.luaRouter.GetRouteStats()
	stats.Components = utils.ComponentStats{
		RouteCount:    routeCount,
		BytecodeBytes: bytecodeBytes,
		SessionStats:  sessions.GlobalSessionManager.GetCacheStats(),
	}
	ctx.SetContentType("text/html; charset=utf-8")
	ctx.SetStatusCode(fasthttp.StatusOK)
	ctx.SetBody([]byte(utils.DebugStatsPage(stats)))
}

// SetStaticCaching enables/disables static file caching
func (s *Server) SetStaticCaching(duration time.Duration) {
	if s.staticFS != nil {
		s.staticFS.CacheDuration = duration
		s.staticHandler = s.staticFS.NewRequestHandler()
	}
}

// GetStaticPrefix returns the URL prefix for static files
func (s *Server) GetStaticPrefix() string {
	return s.staticPrefix
}

// UpdateErrorCache refreshes cached error pages
func (s *Server) UpdateErrorCache() {
	s.errorCacheMu.Lock()
	s.cached404 = []byte(utils.NotFoundPage(s.errorConfig, ""))
	s.cached500 = []byte(utils.InternalErrorPage(s.errorConfig, "", "Internal Server Error"))
	s.errorCacheMu.Unlock()
}