Moonshark/core/runner/Http.go

package runner

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/goccy/go-json"
	"github.com/valyala/bytebufferpool"
	"github.com/valyala/fasthttp"

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

// 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)
}

// ---------- HTTP CLIENT FUNCTIONALITY ----------

// Default HTTP client with sensible timeout
var defaultFastClient fasthttp.Client = fasthttp.Client{
	MaxConnsPerHost:               1024,
	MaxIdleConnDuration:           time.Minute,
	ReadTimeout:                   30 * time.Second,
	WriteTimeout:                  30 * time.Second,
	DisableHeaderNamesNormalizing: true,
}

// HTTPClientConfig contains client settings
type HTTPClientConfig struct {
	// Maximum timeout for requests (0 = no limit)
	MaxTimeout time.Duration
	// Default request timeout
	DefaultTimeout time.Duration
	// Maximum response size in bytes (0 = no limit)
	MaxResponseSize int64
	// Whether to allow remote connections
	AllowRemote bool
}

// DefaultHTTPClientConfig provides sensible defaults
var DefaultHTTPClientConfig = HTTPClientConfig{
	MaxTimeout:      60 * time.Second,
	DefaultTimeout:  30 * time.Second,
	MaxResponseSize: 10 * 1024 * 1024, // 10MB
	AllowRemote:     true,
}

// Function name constant to ensure consistency
const httpRequestFuncName = "__http_request"

// httpRequest makes an HTTP request and returns the result to Lua
func httpRequest(state *luajit.State) int {
	// Get method (required)
	if !state.IsString(1) {
		state.PushString("http.client.request: method must be a string")
		return -1
	}
	method := strings.ToUpper(state.ToString(1))

	// Get URL (required)
	if !state.IsString(2) {
		state.PushString("http.client.request: url must be a string")
		return -1
	}
	urlStr := state.ToString(2)

	// Parse URL to check if it's valid and if it's allowed
	parsedURL, err := url.Parse(urlStr)
	if err != nil {
		state.PushString("Invalid URL: " + err.Error())
		return -1
	}

	// Get client configuration from registry (if available)
	var config HTTPClientConfig = DefaultHTTPClientConfig
	state.GetGlobal("__http_client_config")
	if !state.IsNil(-1) {
		if state.IsTable(-1) {
			// Extract max timeout
			state.GetField(-1, "max_timeout")
			if state.IsNumber(-1) {
				config.MaxTimeout = time.Duration(state.ToNumber(-1)) * time.Second
			}
			state.Pop(1)

			// Extract default timeout
			state.GetField(-1, "default_timeout")
			if state.IsNumber(-1) {
				config.DefaultTimeout = time.Duration(state.ToNumber(-1)) * time.Second
			}
			state.Pop(1)

			// Extract max response size
			state.GetField(-1, "max_response_size")
			if state.IsNumber(-1) {
				config.MaxResponseSize = int64(state.ToNumber(-1))
			}
			state.Pop(1)

			// Extract allow remote
			state.GetField(-1, "allow_remote")
			if state.IsBoolean(-1) {
				config.AllowRemote = state.ToBoolean(-1)
			}
			state.Pop(1)
		}
	}
	state.Pop(1)

	// Check if remote connections are allowed
	if !config.AllowRemote && (parsedURL.Hostname() != "localhost" && parsedURL.Hostname() != "127.0.0.1") {
		state.PushString("Remote connections are not allowed")
		return -1
	}

	// Use bytebufferpool for request and response
	req := fasthttp.AcquireRequest()
	resp := fasthttp.AcquireResponse()
	defer fasthttp.ReleaseRequest(req)
	defer fasthttp.ReleaseResponse(resp)

	// Set up request
	req.Header.SetMethod(method)
	req.SetRequestURI(urlStr)
	req.Header.Set("User-Agent", "Moonshark/1.0")

	// Get body (optional)
	if state.GetTop() >= 3 && !state.IsNil(3) {
		if state.IsString(3) {
			// String body
			req.SetBodyString(state.ToString(3))
		} else if state.IsTable(3) {
			// Table body - convert to JSON
			luaTable, err := state.ToTable(3)
			if err != nil {
				state.PushString("Failed to parse body table: " + err.Error())
				return -1
			}

			// Use bytebufferpool for JSON serialization
			buf := bytebufferpool.Get()
			defer bytebufferpool.Put(buf)

			if err := json.NewEncoder(buf).Encode(luaTable); err != nil {
				state.PushString("Failed to convert body to JSON: " + err.Error())
				return -1
			}

			req.SetBody(buf.Bytes())
		} else {
			state.PushString("Body must be a string or table")
			return -1
		}
	}

	// Process options (headers, timeout, etc.)
	timeout := config.DefaultTimeout
	if state.GetTop() >= 4 && !state.IsNil(4) {
		if !state.IsTable(4) {
			state.PushString("Options must be a table")
			return -1
		}

		// Process headers
		state.GetField(4, "headers")
		if state.IsTable(-1) {
			// Iterate through headers
			state.PushNil() // Start iteration
			for state.Next(-2) {
				// Stack now has key at -2 and value at -1
				if state.IsString(-2) && state.IsString(-1) {
					headerName := state.ToString(-2)
					headerValue := state.ToString(-1)
					req.Header.Set(headerName, headerValue)
				}
				state.Pop(1) // Pop value, leave key for next iteration
			}
		}
		state.Pop(1) // Pop headers table

		// Get timeout
		state.GetField(4, "timeout")
		if state.IsNumber(-1) {
			requestTimeout := time.Duration(state.ToNumber(-1)) * time.Second

			// Apply max timeout if configured
			if config.MaxTimeout > 0 && requestTimeout > config.MaxTimeout {
				timeout = config.MaxTimeout
			} else {
				timeout = requestTimeout
			}
		}
		state.Pop(1) // Pop timeout

		// Set content type for POST/PUT if body is present and content-type not manually set
		if (method == "POST" || method == "PUT") && req.Body() != nil && req.Header.Peek("Content-Type") == nil {
			// Check if options specify content type
			state.GetField(4, "content_type")
			if state.IsString(-1) {
				req.Header.Set("Content-Type", state.ToString(-1))
			} else {
				// Default to JSON if body is a table, otherwise plain text
				if state.IsTable(3) {
					req.Header.Set("Content-Type", "application/json")
				} else {
					req.Header.Set("Content-Type", "text/plain")
				}
			}
			state.Pop(1) // Pop content_type
		}

		// Process query parameters
		state.GetField(4, "query")
		if state.IsTable(-1) {
			// Create URL args
			args := req.URI().QueryArgs()

			// Iterate through query params
			state.PushNil() // Start iteration
			for state.Next(-2) {
				// Stack now has key at -2 and value at -1
				if state.IsString(-2) {
					paramName := state.ToString(-2)

					// Handle different value types
					if state.IsString(-1) {
						args.Add(paramName, state.ToString(-1))
					} else if state.IsNumber(-1) {
						args.Add(paramName, strings.TrimRight(strings.TrimRight(
							state.ToString(-1), "0"), "."))
					} else if state.IsBoolean(-1) {
						if state.ToBoolean(-1) {
							args.Add(paramName, "true")
						} else {
							args.Add(paramName, "false")
						}
					}
				}
				state.Pop(1) // Pop value, leave key for next iteration
			}
		}
		state.Pop(1) // Pop query table
	}

	// Create context with timeout
	_, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	// Execute request
	err = defaultFastClient.DoTimeout(req, resp, timeout)
	if err != nil {
		errStr := "Request failed: " + err.Error()
		if errors.Is(err, fasthttp.ErrTimeout) {
			errStr = "Request timed out after " + timeout.String()
		}
		state.PushString(errStr)
		return -1
	}

	// Create response table
	state.NewTable()

	// Set status code
	state.PushNumber(float64(resp.StatusCode()))
	state.SetField(-2, "status")

	// Set status text
	statusText := fasthttp.StatusMessage(resp.StatusCode())
	state.PushString(statusText)
	state.SetField(-2, "status_text")

	// Set body
	var respBody []byte

	// Apply size limits to response
	if config.MaxResponseSize > 0 && int64(len(resp.Body())) > config.MaxResponseSize {
		// Make a limited copy
		respBody = make([]byte, config.MaxResponseSize)
		copy(respBody, resp.Body())
	} else {
		respBody = resp.Body()
	}

	state.PushString(string(respBody))
	state.SetField(-2, "body")

	// Parse body as JSON if content type is application/json
	contentType := string(resp.Header.ContentType())
	if strings.Contains(contentType, "application/json") {
		var jsonData any
		if err := json.Unmarshal(respBody, &jsonData); err == nil {
			if err := state.PushValue(jsonData); err == nil {
				state.SetField(-2, "json")
			}
		}
	}

	// Set headers
	state.NewTable()
	resp.Header.VisitAll(func(key, value []byte) {
		state.PushString(string(value))
		state.SetField(-2, string(key))
	})
	state.SetField(-2, "headers")

	// Create ok field (true if status code is 2xx)
	state.PushBoolean(resp.StatusCode() >= 200 && resp.StatusCode() < 300)
	state.SetField(-2, "ok")

	return 1
}

// HTTPModuleInitFunc returns an initializer function for the HTTP module
func HTTPModuleInitFunc() StateInitFunc {
	return func(state *luajit.State) error {
		// CRITICAL: Register the native Go function first
		// This must be done BEFORE any Lua code that references it
		if err := state.RegisterGoFunction(httpRequestFuncName, httpRequest); err != nil {
			logger.Error("[HTTP Module] Failed to register __http_request function")
			logger.ErrorCont("%v", err)
			return err
		}

		// Set up default HTTP client configuration
		setupHTTPClientConfig(state)

		// Initialize Lua HTTP module
		if err := state.DoString(LuaHTTPModule); err != nil {
			logger.Error("[HTTP Module] Failed to initialize HTTP module Lua code")
			logger.ErrorCont("%v", err)
			return err
		}

		// Verify HTTP client functions are available
		verifyHTTPClient(state)

		return nil
	}
}

// Helper to set up HTTP client config
func setupHTTPClientConfig(state *luajit.State) {
	state.NewTable()

	state.PushNumber(float64(DefaultHTTPClientConfig.MaxTimeout / time.Second))
	state.SetField(-2, "max_timeout")

	state.PushNumber(float64(DefaultHTTPClientConfig.DefaultTimeout / time.Second))
	state.SetField(-2, "default_timeout")

	state.PushNumber(float64(DefaultHTTPClientConfig.MaxResponseSize))
	state.SetField(-2, "max_response_size")

	state.PushBoolean(DefaultHTTPClientConfig.AllowRemote)
	state.SetField(-2, "allow_remote")

	state.SetGlobal("__http_client_config")
}

// 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
}

// ApplyHTTPResponse applies an HTTP response to a fasthttp.RequestCtx
func ApplyHTTPResponse(httpResp *HTTPResponse, ctx *fasthttp.RequestCtx) {
	// Set status code
	ctx.SetStatusCode(httpResp.Status)

	// Set headers
	for name, value := range httpResp.Headers {
		ctx.Response.Header.Set(name, value)
	}

	// Set cookies
	for _, cookie := range httpResp.Cookies {
		// Convert net/http cookie to fasthttp cookie
		var c fasthttp.Cookie
		c.SetKey(cookie.Name)
		c.SetValue(cookie.Value)
		if cookie.Path != "" {
			c.SetPath(cookie.Path)
		}
		if cookie.Domain != "" {
			c.SetDomain(cookie.Domain)
		}
		if cookie.MaxAge > 0 {
			c.SetMaxAge(cookie.MaxAge)
		}
		if cookie.Expires.After(time.Time{}) {
			c.SetExpire(cookie.Expires)
		}
		c.SetSecure(cookie.Secure)
		c.SetHTTPOnly(cookie.HttpOnly)
		ctx.Response.Header.SetCookie(&c)
	}

	// Process the body based on its type
	if httpResp.Body == nil {
		return
	}

	// Set body based on type
	switch body := httpResp.Body.(type) {
	case string:
		ctx.SetBodyString(body)
	case []byte:
		ctx.SetBody(body)
	case map[string]any, []any, []float64, []string, []int:
		// Marshal JSON using a buffer from the pool
		buf := bytebufferpool.Get()
		defer bytebufferpool.Put(buf)

		if err := json.NewEncoder(buf).Encode(body); err == nil {
			// Set content type if not already set
			if len(ctx.Response.Header.ContentType()) == 0 {
				ctx.Response.Header.SetContentType("application/json")
			}
			ctx.SetBody(buf.Bytes())
		} else {
			// Fallback
			ctx.SetBodyString(fmt.Sprintf("%v", body))
		}
	default:
		// Default to string representation
		ctx.SetBodyString(fmt.Sprintf("%v", body))
	}
}

// WithHTTPClientConfig creates a runner option to configure the HTTP client
func WithHTTPClientConfig(config HTTPClientConfig) RunnerOption {
	return func(r *Runner) {
		// Store the config to be applied during initialization
		r.AddModule("__http_client_config", map[string]any{
			"max_timeout":       float64(config.MaxTimeout / time.Second),
			"default_timeout":   float64(config.DefaultTimeout / time.Second),
			"max_response_size": float64(config.MaxResponseSize),
			"allow_remote":      config.AllowRemote,
		})
	}
}

// RestrictHTTPToLocalhost is a convenience function to restrict HTTP client
// to localhost connections only
func RestrictHTTPToLocalhost() RunnerOption {
	return WithHTTPClientConfig(HTTPClientConfig{
		MaxTimeout:      DefaultHTTPClientConfig.MaxTimeout,
		DefaultTimeout:  DefaultHTTPClientConfig.DefaultTimeout,
		MaxResponseSize: DefaultHTTPClientConfig.MaxResponseSize,
		AllowRemote:     false,
	})
}

// Verify that HTTP client is properly set up
func verifyHTTPClient(state *luajit.State) {
	// Get the client table
	state.GetGlobal("http")
	if !state.IsTable(-1) {
		logger.Warning("[HTTP Module] 'http' is not a table")
		state.Pop(1)
		return
	}

	state.GetField(-1, "client")
	if !state.IsTable(-1) {
		logger.Warning("[HTTP Module] 'http.client' is not a table")
		state.Pop(2)
		return
	}

	// Check for get function
	state.GetField(-1, "get")
	if !state.IsFunction(-1) {
		logger.Warning("[HTTP Module] 'http.client.get' is not a function")
	}
	state.Pop(1)

	// Check for the request function
	state.GetField(-1, "request")
	if !state.IsFunction(-1) {
		logger.Warning("[HTTP Module] 'http.client.request' is not a function")
	}
	state.Pop(3) // Pop request, client, http
}

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
  set_content_type = function(content_type)
      http.set_header("Content-Type", content_type)
  end,

  -- HTTP client submodule
  client = {
    -- Generic request function
    request = function(method, url, body, options)
        if type(method) ~= "string" then
            error("http.client.request: method must be a string", 2)
        end
        if type(url) ~= "string" then
            error("http.client.request: url must be a string", 2)
        end

        -- Call native implementation (this is the critical part)
        local result = __http_request(method, url, body, options)
        return result
    end,

    -- Simple GET request
    get = function(url, options)
        return http.client.request("GET", url, nil, options)
    end,

    -- Simple POST request with automatic content-type
    post = function(url, body, options)
        options = options or {}
        return http.client.request("POST", url, body, options)
    end,

    -- Simple PUT request with automatic content-type
    put = function(url, body, options)
        options = options or {}
        return http.client.request("PUT", url, body, options)
    end,

    -- Simple DELETE request
    delete = function(url, options)
        return http.client.request("DELETE", url, nil, options)
    end,

    -- Simple PATCH request
    patch = function(url, body, options)
        options = options or {}
        return http.client.request("PATCH", url, body, options)
    end,

    -- Simple HEAD request
    head = function(url, options)
        options = options or {}
        local old_options = options
        options = {headers = old_options.headers, timeout = old_options.timeout, query = old_options.query}
        local response = http.client.request("HEAD", url, nil, options)
        return response
    end,

    -- Simple OPTIONS request
    options = function(url, options)
        return http.client.request("OPTIONS", url, nil, options)
    end,

    -- Shorthand function to directly get JSON
    get_json = function(url, options)
        options = options or {}
        local response = http.client.get(url, options)
        if response.ok and response.json then
            return response.json
        end
        return nil, response
    end,

    -- Utility to build a URL with query parameters
    build_url = function(base_url, params)
        if not params or type(params) ~= "table" then
            return base_url
        end

        local query = {}
        for k, v in pairs(params) do
            if type(v) == "table" then
                for _, item in ipairs(v) do
                    table.insert(query, k .. "=" .. tostring(item))
                end
            else
                table.insert(query, k .. "=" .. tostring(v))
            end
        end

        if #query > 0 then
            if base_url:find("?") then
                return base_url .. "&" .. table.concat(query, "&")
            else
                return base_url .. "?" .. table.concat(query, "&")
            end
        end

        return base_url
    end
  }
}

-- Install HTTP module
_G.http = http

-- Clear previous responses when executing scripts
local old_execute_script = __execute_script
if old_execute_script then
  __execute_script = function(fn, ctx)
    -- Clear previous response
    __http_responses[1] = nil

    -- Execute original function
    return old_execute_script(fn, ctx)
  end
end
`