javascript.go

package main

import (
	"fmt"
	"io"
	"log"
	"math/rand"
	"os"
	"sync"

	"github.com/dop251/goja"
	"github.com/envoyproxy/envoy/source/extensions/dynamic_modules/sdk/go/shared"
)

const (
	javaScriptExportedSymbolOnConfig          = "OnConfigure"
	javaScriptExportedSymbolOnRequestHeaders  = "OnRequestHeaders"
	javaScriptExportedSymbolOnResponseHeaders = "OnResponseHeaders"

	functionDeclTemplate = `globalThis.%[1]s = %[1]s`
	numberOfVMPool       = 24
)

type (
	// javaScriptFilterConfigFactory implements [shared.HttpFilterConfigFactory].
	javaScriptFilterConfigFactory struct {
		shared.EmptyHttpFilterConfigFactory
	}
	// javaScriptFilterFactory implements [shared.HttpFilterFactory].
	javaScriptFilterFactory struct {
		vms [numberOfVMPool]*javaScriptVM
	}
	// javaScriptFilter implements [shared.HttpFilter].
	javaScriptFilter struct {
		handle          shared.HttpFilterHandle
		vm              *javaScriptVM
		requestHeaders  map[string]string
		responseHeaders map[string]string
		shared.EmptyHttpFilter
	}
	javaScriptVM struct {
		*goja.Runtime
		mux               sync.Mutex
		onRequestHeaders  goja.Callable
		onResponseHeaders goja.Callable
	}
)

// Create implements [shared.HttpFilterConfigFactory].
func (p *javaScriptFilterConfigFactory) Create(handle shared.HttpFilterConfigHandle, unparsedConfig []byte) (shared.HttpFilterFactory, error) {
	c := &javaScriptFilterFactory{}

	for i := range numberOfVMPool {
		vm, err := newJavaScriptVM(string(unparsedConfig), os.Stdout)
		if err != nil {
			log.Printf("failed to create JavaScript VM: %v", err)
			return nil, err
		}
		c.vms[i] = vm
	}
	return c, nil
}

// Create implements [shared.HttpFilterFactory].
func (p *javaScriptFilterFactory) Create(handle shared.HttpFilterHandle) shared.HttpFilter {
	vm := p.vms[rand.Intn(numberOfVMPool)]
	return &javaScriptFilter{
		handle:          handle,
		vm:              vm,
		requestHeaders:  make(map[string]string),
		responseHeaders: make(map[string]string),
	}
}

func newJavaScriptVM(script string, w io.Writer) (*javaScriptVM, error) {
	vm := goja.New()
	console := vm.NewObject()
	err := console.Set("log", func(call goja.FunctionCall) goja.Value {
		args := make([]interface{}, 0, len(call.Arguments))
		for _, a := range call.Arguments {
			args = append(args, a.Export())
		}
		_, _ = fmt.Fprint(w, args...)
		return goja.Undefined()
	})
	if err != nil {
		return nil, fmt.Errorf("failed to set console.log: %w", err)
	}
	err = vm.Set("console", console)
	if err != nil {
		return nil, fmt.Errorf("failed to set console: %w", err)
	}

	_, err = vm.RunString(script)
	if err != nil {
		return nil, fmt.Errorf("failed to run script: %w", err)
	}

	// Call OnConfigure.
	onConfigure, ok := goja.AssertFunction(vm.GlobalObject().Get(javaScriptExportedSymbolOnConfig))
	if !ok {
		return nil, fmt.Errorf("failed to get %s function", javaScriptExportedSymbolOnConfig)
	}
	_, err = onConfigure(goja.Undefined())
	if err != nil {
		return nil, fmt.Errorf("failed to call %s function: %w", javaScriptExportedSymbolOnConfig, err)
	}

	ret := &javaScriptVM{Runtime: vm}
	// Check two exported functions.
	ret.onRequestHeaders, ok = goja.AssertFunction(vm.GlobalObject().Get(javaScriptExportedSymbolOnRequestHeaders))
	if !ok {
		return nil, fmt.Errorf("failed to get %s function", javaScriptExportedSymbolOnRequestHeaders)
	}
	ret.onResponseHeaders, ok = goja.AssertFunction(vm.GlobalObject().Get(javaScriptExportedSymbolOnResponseHeaders))
	if !ok {
		return nil, fmt.Errorf("failed to get %s function", javaScriptExportedSymbolOnResponseHeaders)
	}
	return ret, nil
}

// OnRequestHeaders implements [shared.HttpFilter].
func (p *javaScriptFilter) OnRequestHeaders(headers shared.HeaderMap, _ bool) shared.HeadersStatus {
	for _, header := range headers.GetAll() {
		p.requestHeaders[header[0]] = header[1]
	}
	p.vm.mux.Lock()
	defer p.vm.mux.Unlock()
	vm := p.vm
	obj := vm.NewObject()
	_ = obj.Set("getRequestHeader", func(call goja.FunctionCall) goja.Value {
		if len(call.Arguments) < 1 {
			return vm.ToValue("")
		}
		key := call.Argument(0).String()
		return vm.ToValue(p.requestHeaders[key])
	})
	_ = obj.Set("setRequestHeader", func(call goja.FunctionCall) goja.Value {
		if len(call.Arguments) < 2 {
			return goja.Undefined()
		}
		key := call.Argument(0).String()
		value := call.Argument(1).String()
		p.requestHeaders[key] = value
		headers.Set(key, value)
		return goja.Undefined()
	})
	if _, err := vm.onRequestHeaders(goja.Undefined(), obj); err != nil {
		log.Printf("failed to call %s: %v", javaScriptExportedSymbolOnRequestHeaders, err)
		return shared.HeadersStatusStop
	}
	return shared.HeadersStatusContinue
}

// OnResponseHeaders implements [shared.HttpFilter].
func (p *javaScriptFilter) OnResponseHeaders(headers shared.HeaderMap, _ bool) shared.HeadersStatus {
	for _, header := range headers.GetAll() {
		p.responseHeaders[header[0]] = header[1]
	}
	p.vm.mux.Lock()
	defer p.vm.mux.Unlock()
	vm := p.vm
	obj := vm.NewObject()
	_ = obj.Set("getRequestHeader", func(call goja.FunctionCall) goja.Value {
		if len(call.Arguments) < 1 {
			return vm.ToValue("")
		}
		key := call.Argument(0).String()
		return vm.ToValue(p.requestHeaders[key])
	})

	// Setting request header in response phase is not allowed.

	_ = obj.Set("getResponseHeader", func(call goja.FunctionCall) goja.Value {
		if len(call.Arguments) < 1 {
			return vm.ToValue("")
		}
		key := call.Argument(0).String()
		return vm.ToValue(p.responseHeaders[key])
	})
	_ = obj.Set("setResponseHeader", func(call goja.FunctionCall) goja.Value {
		if len(call.Arguments) < 2 {
			return goja.Undefined()
		}
		key := call.Argument(0).String()
		value := call.Argument(1).String()
		p.responseHeaders[key] = value
		headers.Set(key, value)
		return goja.Undefined()
	})
	if _, err := vm.onResponseHeaders(goja.Undefined(), obj); err != nil {
		log.Printf("failed to call %s: %v", javaScriptExportedSymbolOnResponseHeaders, err)
		return shared.HeadersStatusStop
	}
	return shared.HeadersStatusContinue
}