buildr/modules/plugin/exec_context.go

package plugin

import (
	"bytes"
	"context"
	"errors"
	"log/slog"
	"os/exec"
	"reflect"
	"time"

	"github.com/tetratelabs/wazero/api"

	bapi "code.icb4dc0.de/buildr/api"
	remotev1 "code.icb4dc0.de/buildr/api/generated/remote/v1"
	wasiv1 "code.icb4dc0.de/buildr/api/generated/wasi/v1"

	"code.icb4dc0.de/buildr/buildr/internal/sh"
	"code.icb4dc0.de/buildr/buildr/modules"
)

const defaultIntegerSize = 32

func NewHostFuncExport[TRequest bapi.ProtoMessage, TResponse bapi.ProtoMessage](f func(ctx context.Context, req TRequest) (TResponse, error)) HostFuncExportWrapper[TRequest, TResponse] {
	return f
}

type HostFuncExportWrapper[TRequest bapi.ProtoMessage, TResponse bapi.ProtoMessage] func(ctx context.Context, req TRequest) (TResponse, error)

func (w HostFuncExportWrapper[TRequest, TResponse]) Call(ctx context.Context, m api.Module, offset, byteCount uint32) uint64 {
	buf, ok := m.Memory().Read(offset, byteCount)
	if !ok {
		return 0
	}

	var req TRequest
	if reflect.TypeOf(req).Kind() == reflect.Ptr {
		req = reflect.New(reflect.TypeOf(req).Elem()).Interface().(TRequest)
	}

	if err := req.UnmarshalVT(buf); err != nil {
		panic(err)
	}

	resp, err := w(ctx, req)
	if err != nil {
		panic(err)
	}

	data, err := resp.MarshalVT()
	if err != nil {
		panic(err)
	}

	ptr, err := w.allocate(ctx, m, uint64(len(data)))
	if err != nil {
		panic(err)
	}

	if !m.Memory().Write(uint32(ptr), data) {
		panic("failed to write message to memory")
	}

	return (ptr << uint64(defaultIntegerSize)) | uint64(len(data))
}

func (w HostFuncExportWrapper[TRequest, TResponse]) allocate(ctx context.Context, m api.Module, size uint64) (ptr uint64, err error) {
	results, err := m.ExportedFunction("allocate").Call(ctx, size)
	if err != nil {
		return 0, err
	}

	return results[0], nil
}

type wasiPluginExecContext struct {
	ctx modules.ExecutionContext
}

func (c *wasiPluginExecContext) lookPath(_ context.Context, lookupPathReq *wasiv1.LookupPathRequest) (*wasiv1.LookupPathResponse, error) {
	if lookupPathReq.Command == "" {
		return nil, errors.New("cannot lookup empty command")
	}

	foundPath, err := exec.LookPath(lookupPathReq.Command)

	lookupPathResp := &wasiv1.LookupPathResponse{
		Path: foundPath,
	}

	if err != nil {
		lookupPathResp.Error = err.Error()
	}

	return lookupPathResp, nil
}

func (c *wasiPluginExecContext) exec(_ context.Context, execReq *wasiv1.ProcessStartRequest) (*wasiv1.ProcessStartResponse, error) {
	execResp := new(wasiv1.ProcessStartResponse)

	//nolint:contextcheck // that ain't a context.Context
	cmd, err := sh.PrepareCommand(c.ctx, execReq.Command, execReq.Args...)
	if err != nil {
		c.ctx.Logger().Error("failed to prepare command", slog.String("err", err.Error()))
		return nil, err
	}

	cmd.AddEnv(execReq.Environment)
	if len(execReq.Stdin) > 0 {
		cmd.SetStdIn(bytes.NewReader(execReq.Stdin))
	}

	if len(execReq.WorkingDirectory) > 0 {
		cmd.SetWorkingDir(execReq.WorkingDirectory)
	}

	if err := cmd.Run(); err != nil {
		var exitErr *exec.ExitError
		if errors.As(err, &exitErr) {
			execResp.ExitCode = int32(exitErr.ExitCode())
			execResp.Stderr = exitErr.Stderr
		} else {
			execResp.ExitCode = 1
		}
		execResp.Error = err.Error()
	}

	return execResp, nil
}

func (c *wasiPluginExecContext) log(ctx context.Context, taskLog *remotev1.TaskLog) (*remotev1.Result, error) {
	rec := slog.NewRecord(time.UnixMicro(taskLog.Time), slog.Level(taskLog.Level), taskLog.Message, 0)

	for i := range taskLog.Attributes {
		attr := taskLog.Attributes[i]
		rec.AddAttrs(slog.String(attr.Key, attr.Value))
	}

	if err := c.ctx.Logger().Handler().Handle(ctx, rec); err != nil {
		return &remotev1.Result{Error: err.Error()}, nil
	}

	return new(remotev1.Result), nil
}

func (c *wasiPluginExecContext) getState(ctx context.Context, getStateReq *remotev1.GetStateRequest) (*remotev1.GetStateResponse, error) {
	resp := new(remotev1.GetStateResponse)
	val, _, err := c.ctx.GetState(ctx, string(getStateReq.Key))
	if err != nil {
		c.ctx.Logger().Error("failed to get state", slog.String("err", err.Error()))
		return nil, err
	}

	resp.Data = val

	return resp, nil
}

func (c *wasiPluginExecContext) setState(ctx context.Context, setState *remotev1.SetState) (*remotev1.Result, error) {
	var resp remotev1.Result
	if err := c.ctx.SetState(ctx, string(setState.Key), setState.Data); err != nil {
		c.ctx.Logger().Error("failed to set state", slog.String("err", err.Error()))
		resp.Error = err.Error()
	}

	return &resp, nil
}