Moonshark/core/runner/Http.go

package runner

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"fmt"
	"net/url"
	"strings"
	"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"
)

// Default HTTP client with sensible timeout
var defaultFastClient = 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,
}

// ApplyResponse applies a Response to a fasthttp.RequestCtx
func ApplyResponse(resp *Response, ctx *fasthttp.RequestCtx) {
	// Set status code
	ctx.SetStatusCode(resp.Status)

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

	// Set cookies
	for _, cookie := range resp.Cookies {
		ctx.Response.Header.SetCookie(cookie)
	}

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

	// Get a buffer from the pool
	buf := bytebufferpool.Get()
	defer bytebufferpool.Put(buf)

	// Set body based on type
	switch body := resp.Body.(type) {
	case string:
		ctx.SetBodyString(body)
	case []byte:
		ctx.SetBody(body)
	case map[string]any, []any, []float64, []string, []int:
		// Marshal JSON
		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))
	}
}

// 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
	parsedURL, err := url.Parse(urlStr)
	if err != nil {
		state.PushString("Invalid URL: " + err.Error())
		return -1
	}

	// Get client configuration
	config := DefaultHTTPClientConfig

	// 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())
			req.Header.SetContentType("application/json")
		} 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) && state.IsTable(4) {
		// 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

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

// generateToken creates a cryptographically secure random token
func generateToken(state *luajit.State) int {
	// Get the length from the Lua arguments (default to 32)
	length := 32
	if state.GetTop() >= 1 && state.IsNumber(1) {
		length = int(state.ToNumber(1))
	}

	// Enforce minimum length for security
	if length < 16 {
		length = 16
	}

	// Generate secure random bytes
	tokenBytes := make([]byte, length)
	if _, err := rand.Read(tokenBytes); err != nil {
		logger.Error("Failed to generate secure token: %v", err)
		state.PushString("")
		return 1 // Return empty string on error
	}

	// Encode as base64
	token := base64.RawURLEncoding.EncodeToString(tokenBytes)

	// Trim to requested length (base64 might be longer)
	if len(token) > length {
		token = token[:length]
	}

	// Push the token to the Lua stack
	state.PushString(token)
	return 1 // One return value
}