functions

cmd_function.go (6014B)

---

 package main
 
 import (
 	"context"
 	"errors"
 	"fmt"
 	"log/slog"
 	"net/http"
 	"os"
 	"path"
 
 	"github.com/firecracker-microvm/firecracker-go-sdk"
 	"github.com/valyala/bytebufferpool"
 	"go.lair.cx/functions/configs"
 	"go.lair.cx/functions/internal/functions"
 	"go.lair.cx/functions/internal/vmutils"
 )
 
 const (
 	DefaultRootfsName = "rootfs.ext4"
 )
 
 type FunctionCmd struct {
 	List   FunctionListCmd   `cmd:"" help:"List functions."`
 	Create FunctionCreateCmd `cmd:"" help:"Create a function."`
 	Run    FunctionRunCmd    `cmd:"" help:"Run a function."`
 }
 
 type FunctionListCmd struct {
 	Namespace string `short:"N" arg:"" help:"Namespace of the function."`
 }
 
 func (c *FunctionListCmd) Run() error {
 	namespacePath := path.Join(configs.FuncDir, c.Namespace)
 
 	stat, err := os.Stat(namespacePath)
 	if err != nil {
 		return err
 	}
 
 	if !stat.IsDir() {
 		return errors.New("namespace is not a directory")
 	}
 
 	files, err := os.ReadDir(namespacePath)
 	if err != nil {
 		return err
 	}
 
 	for _, file := range files {
 		if !checkFilesInDir(
 			path.Join(namespacePath, file.Name()),
 			"function.json",
 		) {
 			continue
 		}
 
 		fmt.Println(file.Name())
 	}
 
 	return nil
 }
 
 type FunctionCreateCmd struct {
 	Namespace string `short:"N" arg:"" help:"Namespace of the function."`
 	Name      string `short:"n" arg:"" help:"Name of the function."`
 
 	Kernel string `short:"k" required:"" help:"File path for kernal image."`
 	Rootfs string `short:"r" help:"File path for rootfs image. Can be omitted if --source is provided."`
 	Source string `short:"s" help:"File path for source code. Can be omitted if --rootfs is provided."`
 	NoCopy bool   `help:"Instead of copying the file to the function folder, use that path as is."`
 }
 
 func (c *FunctionCreateCmd) Run() error {
 	namespacePath := path.Join(configs.FuncDir, c.Namespace)
 
 	stat, err := os.Stat(namespacePath)
 	if err != nil {
 		return err
 	}
 
 	if !stat.IsDir() {
 		return errors.New("namespace is not a directory")
 	}
 
 	functionPath := path.Join(namespacePath, c.Name)
 
 	if _, err := os.Stat(functionPath); err == nil {
 		return errors.New("function already exists")
 	}
 
 	err = os.Mkdir(functionPath, 0o755)
 	if err != nil {
 		return err
 	}
 
 	// --source or --rootfs must be provided.
 	if len(c.Source) == 0 && len(c.Rootfs) == 0 {
 		return errors.New("either --source or --rootfs must be provided")
 	}
 
 	// When --rootfs is not provided, we use the default name.
 	if len(c.Rootfs) == 0 {
 		c.Rootfs = path.Join(functionPath, DefaultRootfsName)
 	}
 
 	// Trying to build image from source, if provided.
 	if len(c.Source) > 0 {
 		if _, err := os.Stat(c.Rootfs); err == nil {
 			return errors.New("rootfs already exists")
 		}
 
 		err = functions.BuildRootfsFromSource(c.Source, c.Rootfs)
 		if err != nil {
 			return err
 		}
 	}
 
 	// When no-copy mode, the kernel and rootfs path must be absolute.
 	if c.NoCopy {
 		if !path.IsAbs(c.Kernel) {
 			return errors.New("kernel path must be absolute")
 		}
 
 		if !path.IsAbs(c.Rootfs) {
 			return errors.New("rootfs path must be absolute")
 		}
 	} else {
 		/*
 			kernelPath, err = copyFileToDir(c.Kernel, functionPath)
 			if err != nil {
 				return err
 			}
 
 			if len(c.Rootfs) > 0 {
 				rootfsPath, err = copyFileToDir(c.Rootfs, functionPath)
 				if err != nil {
 					return err
 				}
 			}
 		*/
 	}
 
 	return nil
 }
 
 type FunctionRunCmd struct {
 	Namespace string `short:"N" arg:"" help:"Namespace of the function."`
 	Name      string `short:"n" arg:"" help:"Name of the function."`
 	Kernel    string `short:"k" required:"" help:"File path for kernal image."`
 	Network   string `short:"n" help:"CNI network to use."`
 	Vsock     string `short:"V" required:"" help:"Path to vsock file."`
 
 	Attach  bool `short:"a" help:"Attach to the function vm."`
 	Verbose bool `short:"v" help:"Enable verbose output."`
 
 	Headers []string `name:"header" short:"H" help:"Extra headers to include in request."`
 	Data    string   `short:"d" help:"Data to pass to the function."`
 }
 
 func (c *FunctionRunCmd) Run() error {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
 	machine, err := c.buildMachine()
 	if err != nil {
 		return fmt.Errorf("failed to build machine: %w", err)
 	}
 
 	err = machine.Start(ctx)
 	if err != nil {
 		return fmt.Errorf("failed to start machine: %w", err)
 	}
 
 	defer c.cleanup(machine)
 
 	client, err := vmutils.NewHTTPClient(machine)
 	if err != nil {
 		return fmt.Errorf("failed to create http client: %w", err)
 	}
 
 	// TODO: Setup network if provided.
 	if len(c.Network) > 0 {
 		ip, gw, err := vmutils.GetNetworkInfo(machine)
 		if err != nil {
 			slog.Error("failed to get network info", "error", err)
 		} else {
 			slog.Info("network info", "ip", ip, "gw", gw)
 		}
 	}
 
 	buf := bytebufferpool.Get()
 	defer bytebufferpool.Put(buf)
 
 	err = runFunction(
 		client,
 		http.Header{},
 		[]byte(c.Data),
 		func(resp *http.Response) (err error) {
 			buf.ReadFrom(resp.Body)
 			return
 		},
 	)
 	if err != nil {
 		return err
 	}
 
 	fmt.Print(buf.String())
 
 	err = machine.Wait(ctx)
 	if err != nil {
 		return fmt.Errorf("vmutils: failed to wait for machine: %w", err)
 	}
 
 	return nil
 }
 
 func (c *FunctionRunCmd) buildMachine() (*firecracker.Machine, error) {
 	rootfsPath, err := getRootfsPath(c.Namespace, c.Name)
 	if err != nil {
 		return nil, err
 	}
 
 	configs := []func(*vmutils.MachineBuilder){
 		vmutils.WithVsock(c.Vsock),
 		vmutils.WithVerbose(c.Verbose),
 	}
 
 	if c.Attach {
 		configs = append(
 			configs,
 			vmutils.WithStdin(os.Stdin),
 			vmutils.WithStdout(os.Stdout),
 			vmutils.WithStderr(os.Stderr),
 		)
 	}
 
 	if c.Network != "" {
 		configs = append(configs, vmutils.WithNetwork(c.Network))
 	}
 
 	return vmutils.BuildMachine(
 		c.Kernel,
 		rootfsPath,
 		configs...,
 	)
 }
 
 func (c *FunctionRunCmd) cleanup(machine *firecracker.Machine) error {
 	// Stop the VMM to ensure that the VM is stopped
 	err := machine.StopVMM()
 	if err != nil {
 		return fmt.Errorf("failed to stop VMM: %w", err)
 	}
 
 	os.Remove(machine.Cfg.SocketPath)
 
 	for _, vsock := range machine.Cfg.VsockDevices {
 		os.Remove(vsock.Path)
 	}
 
 	return nil
 }