Moonshark/core/runner/http.go

package runner

import (
	"net/http"
	"sync"

	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
)

// HTTPResponse represents an HTTP response from Lua
type HTTPResponse struct {
	Status  int               `json:"status"`
	Headers map[string]string `json:"headers"`
	Body    any               `json:"body"`
	Cookies []*http.Cookie    `json:"-"`
}

// Response pool to reduce allocations
var responsePool = sync.Pool{
	New: func() interface{} {
		return &HTTPResponse{
			Status:  200,
			Headers: make(map[string]string, 8), // Pre-allocate with reasonable capacity
			Cookies: make([]*http.Cookie, 0, 4), // Pre-allocate with reasonable capacity
		}
	},
}

// NewHTTPResponse creates a default HTTP response, potentially reusing one from the pool
func NewHTTPResponse() *HTTPResponse {
	return responsePool.Get().(*HTTPResponse)
}

// ReleaseResponse returns the response to the pool after clearing its values
func ReleaseResponse(resp *HTTPResponse) {
	if resp == nil {
		return
	}

	// Clear all values to prevent data leakage
	resp.Status = 200 // Reset to default

	// Clear headers
	for k := range resp.Headers {
		delete(resp.Headers, k)
	}

	// Clear cookies
	resp.Cookies = resp.Cookies[:0] // Keep capacity but set length to 0

	// Clear body
	resp.Body = nil

	responsePool.Put(resp)
}

// LuaHTTPModule is the pure Lua implementation of the HTTP module
const LuaHTTPModule = `
-- Table to store response data
__http_responses = {}

-- HTTP module implementation
local http = {
  -- Set HTTP status code
  set_status = function(code)
    if type(code) ~= "number" then
      error("http.set_status: status code must be a number", 2)
    end

    local resp = __http_responses[1] or {}
    resp.status = code
    __http_responses[1] = resp
  end,

  -- Set HTTP header
  set_header = function(name, value)
    if type(name) ~= "string" or type(value) ~= "string" then
      error("http.set_header: name and value must be strings", 2)
    end

    local resp = __http_responses[1] or {}
    resp.headers = resp.headers or {}
    resp.headers[name] = value
    __http_responses[1] = resp
  end
}

-- Set content type; set_header helper
http.set_content_type = function(content_type)
    http.set_header("Content-Type", content_type)
end

-- Install HTTP module
_G.http = http

-- Override sandbox executor to clear HTTP responses
local old_execute_sandbox = __execute_sandbox
__execute_sandbox = function(bytecode, ctx)
  -- Clear previous response for this thread
  __http_responses[1] = nil

  -- Execute the original function
  local result = old_execute_sandbox(bytecode, ctx)

  -- Return the result unchanged
  return result
end

-- Make sure the HTTP module is accessible in sandbox
if __env_system and __env_system.base_env then
  __env_system.base_env.http = http
end
`

// HTTPModuleInitFunc returns an initializer function for the HTTP module
func HTTPModuleInitFunc() StateInitFunc {
	return func(state *luajit.State) error {
		// Initialize pure Lua HTTP module
		return state.DoString(LuaHTTPModule)
	}
}

// GetHTTPResponse extracts the HTTP response from Lua state
func GetHTTPResponse(state *luajit.State) (*HTTPResponse, bool) {
	response := NewHTTPResponse()

	// Get response table
	state.GetGlobal("__http_responses")
	if state.IsNil(-1) {
		state.Pop(1)
		ReleaseResponse(response) // Return unused response to pool
		return nil, false
	}

	// Check for response at thread index
	state.PushNumber(1)
	state.GetTable(-2)
	if state.IsNil(-1) {
		state.Pop(2)
		ReleaseResponse(response) // Return unused response to pool
		return nil, false
	}

	// Get status
	state.GetField(-1, "status")
	if state.IsNumber(-1) {
		response.Status = int(state.ToNumber(-1))
	}
	state.Pop(1)

	// Get headers
	state.GetField(-1, "headers")
	if state.IsTable(-1) {
		// Iterate through headers table
		state.PushNil() // Start iteration
		for state.Next(-2) {
			// Stack has key at -2 and value at -1
			if state.IsString(-2) && state.IsString(-1) {
				key := state.ToString(-2)
				value := state.ToString(-1)
				response.Headers[key] = value
			}
			state.Pop(1) // Pop value, leave key for next iteration
		}
	}
	state.Pop(1)

	// Get cookies
	state.GetField(-1, "cookies")
	if state.IsTable(-1) {
		// Iterate through cookies array
		length := state.GetTableLength(-1)
		for i := 1; i <= length; i++ {
			state.PushNumber(float64(i))
			state.GetTable(-2)

			if state.IsTable(-1) {
				cookie := extractCookie(state)
				if cookie != nil {
					response.Cookies = append(response.Cookies, cookie)
				}
			}
			state.Pop(1)
		}
	}
	state.Pop(1)

	// Clean up
	state.Pop(2) // Pop response table and __http_responses

	return response, true
}