Moonshark/core/runner/sandbox/Http.go

package sandbox

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

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

	"Moonshark/core/utils/logger"

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

// SessionHandler interface for session management
type SessionHandler interface {
	LoadSession(ctx *fasthttp.RequestCtx) (string, map[string]any)
	SaveSession(ctx *fasthttp.RequestCtx, sessionID string, data map[string]any) bool
}

// 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         []*fasthttp.Cookie `json:"-"`
	SessionModified bool               `json:"-"`
}

// Response pool to reduce allocations
var responsePool = sync.Pool{
	New: func() any {
		return &HTTPResponse{
			Status:  200,
			Headers: make(map[string]string, 8),
			Cookies: make([]*fasthttp.Cookie, 0, 4),
		}
	},
}

// 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 {
	MaxTimeout      time.Duration // Maximum timeout for requests (0 = no limit)
	DefaultTimeout  time.Duration // Default request timeout
	MaxResponseSize int64         // Maximum response size in bytes (0 = no limit)
	AllowRemote     bool          // Whether to allow remote connections
}

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

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

// ReleaseResponse returns the response to the pool
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

	// Reset session flag
	resp.SessionModified = false

	// Clear body
	resp.Body = nil

	responsePool.Put(resp)
}

// HTTPModuleInitFunc returns an initializer function for the HTTP module
func HTTPModuleInitFunc() func(*luajit.State) error {
	return func(state *luajit.State) error {
		// Register the native Go function first
		if err := state.RegisterGoFunction("__http_request", httpRequest); err != nil {
			logger.Error("[HTTP Module] Failed to register __http_request function: %v", err)
			return err
		}

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

		return nil
	}
}

// setupHTTPClientConfig configures HTTP client in Lua
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 nil, false
	}

	// Check for response at thread index
	state.PushNumber(1)
	state.GetTable(-2)
	if state.IsNil(-1) {
		state.Pop(2)
		ReleaseResponse(response)
		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)

	// Check if session was modified
	state.GetGlobal("__session_modified")
	if state.IsBoolean(-1) && state.ToBoolean(-1) {
		response.SessionModified = true
	}
	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 {
		ctx.Response.Header.SetCookie(cookie)
	}

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

// extractCookie grabs cookies from the Lua state
func extractCookie(state *luajit.State) *fasthttp.Cookie {
	cookie := fasthttp.AcquireCookie()

	// Get name
	state.GetField(-1, "name")
	if !state.IsString(-1) {
		state.Pop(1)
		fasthttp.ReleaseCookie(cookie)
		return nil // Name is required
	}
	cookie.SetKey(state.ToString(-1))
	state.Pop(1)

	// Get value
	state.GetField(-1, "value")
	if state.IsString(-1) {
		cookie.SetValue(state.ToString(-1))
	}
	state.Pop(1)

	// Get path
	state.GetField(-1, "path")
	if state.IsString(-1) {
		cookie.SetPath(state.ToString(-1))
	} else {
		cookie.SetPath("/") // Default path
	}
	state.Pop(1)

	// Get domain
	state.GetField(-1, "domain")
	if state.IsString(-1) {
		cookie.SetDomain(state.ToString(-1))
	}
	state.Pop(1)

	// Get expires
	state.GetField(-1, "expires")
	if state.IsNumber(-1) {
		expiry := int64(state.ToNumber(-1))
		cookie.SetExpire(time.Unix(expiry, 0))
	}
	state.Pop(1)

	// Get max age
	state.GetField(-1, "max_age")
	if state.IsNumber(-1) {
		cookie.SetMaxAge(int(state.ToNumber(-1)))
	}
	state.Pop(1)

	// Get secure
	state.GetField(-1, "secure")
	if state.IsBoolean(-1) {
		cookie.SetSecure(state.ToBoolean(-1))
	}
	state.Pop(1)

	// Get http only
	state.GetField(-1, "http_only")
	if state.IsBoolean(-1) {
		cookie.SetHTTPOnly(state.ToBoolean(-1))
	}
	state.Pop(1)

	return cookie
}

// 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
	var config HTTPClientConfig = DefaultHTTPClientConfig
	state.GetGlobal("__http_client_config")
	if !state.IsNil(-1) && 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
}