http-boot-server/bootserver/bootprotocol/bootprotocol.go

// package bootprotocol contains the elements necessary to use the custom network boot protocol
package bootprotocol

import (
	"bytes"
	"encoding"
	"errors"
	"fmt"
	"strings"

	"github.com/google/uuid"
)

type Action int8

const (
	ActionRequest Action = iota
	ActionAccept
	ActionDeny
	ActionDiscover
	ActionUnknown
)

var spaceByte = []byte(" ")
var commaByte = []byte(";")

const (
	keyID     = "id"
	keyEfiApp = "efi_app"
	keyReason = "reason"
)

var ErrInvalidFormat = errors.New("invalid format for message")
var ErrUnknownAction = errors.New("unknown action for message")
var ErrInvalidParam = errors.New("invalid parameter for message")
var ErrMissingParam = errors.New("missing parameter for message")

func (a Action) String() string {
	switch a {
	case ActionAccept:
		return "BOOT_ACCEPT"
	case ActionDeny:
		return "BOOT_DENY"
	case ActionRequest:
		return "BOOT_REQUEST"
	case ActionDiscover:
		return "BOOT_DISCOVER"
	default:
		return "unknown"
	}
}

func newActionFromBytes(raw []byte) Action {
	switch string(raw) {
	case "BOOT_ACCEPT":
		return ActionAccept
	case "BOOT_DENY":
		return ActionDeny
	case "BOOT_REQUEST":
		return ActionRequest
	case "BOOT_DISCOVER":
		return ActionDiscover
	default:
		return ActionUnknown
	}
}

type Message interface {
	encoding.BinaryUnmarshaler
	encoding.BinaryMarshaler
	Action() Action
	ID() uuid.UUID
	String() string
}

type requestMessage struct {
	id uuid.UUID
}

func (rm *requestMessage) UnmarshalBinary(data []byte) error {
	params := bytes.Split(data, commaByte)
	for _, p := range params {
		k, v, err := splitKeyValue(p)
		if err != nil {
			return fmt.Errorf("failed to parse parameter %q: %w", string(p), err)
		}
		if bytes.Equal(k, []byte(keyID)) {
			parsedId, err := uuid.ParseBytes(v)
			if err != nil {
				return ErrInvalidParam
			}
			rm.id = parsedId
			return nil
		}
	}
	return ErrMissingParam
}

func (rm *requestMessage) MarshalBinary() (data []byte, err error) {
	action := []byte(rm.Action().String())
	params := []byte(fmt.Sprintf("%s=%s", keyID, rm.id.String()))
	return bytes.Join([][]byte{action, params}, spaceByte), nil
}

func (rm *requestMessage) Action() Action {
	return ActionRequest
}

func (rm *requestMessage) ID() uuid.UUID {
	return rm.id
}

func (rm *requestMessage) String() string {
	return fmt.Sprintf("%s from %s", ActionRequest.String(), rm.ID().String())
}

type acceptMessage struct {
	id     uuid.UUID
	efiApp string
}

func (am *acceptMessage) UnmarshalBinary(data []byte) error {
	params := bytes.Split(data, commaByte)
	for _, p := range params {
		k, v, err := splitKeyValue(p)
		if err != nil {
			return fmt.Errorf("failed to parse parameter %q: %w", string(p), err)
		}
		switch string(k) {
		case keyID:
			parsedId, err := uuid.ParseBytes(v)
			if err != nil {
				return ErrInvalidParam
			}
			am.id = parsedId
		case keyEfiApp:
			am.efiApp = string(v)
		}
	}

	if am.id == uuid.Nil || am.efiApp == "" {
		return ErrMissingParam
	}
	return nil
}

func (am *acceptMessage) MarshalBinary() (data []byte, err error) {
	action := []byte(am.Action().String())
	// efiApp := strings.ReplaceAll(am.efiApp, `\`, `\\`)
	// efiApp = strings.ReplaceAll(efiApp, "File(", "")
	efiApp := strings.ReplaceAll(am.efiApp, "File(", "")
	efiApp, _ = strings.CutSuffix(efiApp, ")")
	params := [][]byte{
		[]byte(fmt.Sprintf("%s=%s", keyID, am.id.String())),
		[]byte(fmt.Sprintf("%s=%s", keyEfiApp, efiApp)),
	}
	param_bytes := bytes.Join(params, commaByte)
	return bytes.Join([][]byte{action, param_bytes}, spaceByte), nil
}

func (am *acceptMessage) Action() Action {
	return ActionAccept
}

func (am *acceptMessage) ID() uuid.UUID {
	return am.id
}

func (am *acceptMessage) String() string {
	return fmt.Sprintf("%s from %s, app %s", ActionAccept.String(), am.ID().String(), am.efiApp)
}

type denyMessage struct {
	id     uuid.UUID
	reason string
}

func (dm *denyMessage) UnmarshalBinary(data []byte) error {
	params := bytes.Split(data, commaByte)
	for _, p := range params {
		k, v, err := splitKeyValue(p)
		if err != nil {
			return fmt.Errorf("failed to parse parameter %q: %w", string(p), err)
		}
		switch string(k) {
		case keyID:
			parsedId, err := uuid.ParseBytes(v)
			if err != nil {
				return ErrInvalidParam
			}
			dm.id = parsedId
		case keyReason:
			dm.reason = string(v)
		}
	}

	if dm.id == uuid.Nil || dm.reason == "" {
		return ErrMissingParam
	}
	return nil
}

func (dm *denyMessage) MarshalBinary() (data []byte, err error) {
	action := []byte(dm.Action().String())
	params := [][]byte{
		[]byte(fmt.Sprintf("%s=%s", keyID, dm.id.String())),
		[]byte(fmt.Sprintf("%s=%s", keyReason, dm.reason)),
	}
	param_bytes := bytes.Join(params, commaByte)
	return bytes.Join([][]byte{action, param_bytes}, spaceByte), nil
}

func (dm *denyMessage) Action() Action {
	return ActionDeny
}

func (dm *denyMessage) ID() uuid.UUID {
	return dm.id
}

func (dm *denyMessage) String() string {
	return fmt.Sprintf("%s from %s, reason %q", ActionDeny.String(), dm.ID().String(), dm.reason)
}

type discoverMessage struct{}

func (dm *discoverMessage) UnmarshalBinary(data []byte) error {
	return nil
}

func (dm *discoverMessage) MarshalBinary() (data []byte, err error) {
	return []byte(dm.Action().String()), nil
}

func (dm *discoverMessage) Action() Action {
	return ActionDiscover
}

func (dm *discoverMessage) ID() uuid.UUID {
	return uuid.Nil
}

func (dm *discoverMessage) String() string {
	return ActionDiscover.String()
}

func MessageFromBytes(dat []byte) (Message, error) {
	rawAction, content, _ := bytes.Cut(dat, spaceByte)

	var message Message
	action := newActionFromBytes(rawAction)
	switch action {
	case ActionRequest:
		message = &requestMessage{}
	case ActionAccept:
		message = &acceptMessage{}
	case ActionDeny:
		message = &denyMessage{}
	case ActionDiscover:
		message = &discoverMessage{}
	default:
		return nil, ErrUnknownAction
	}

	if err := message.UnmarshalBinary(content); err != nil {
		return nil, fmt.Errorf("failed to parse %s message: %w", message.Action().String(), err)
	}
	return message, nil
}

func Accept(id uuid.UUID, efiApp string) Message {
	return &acceptMessage{
		id:     id,
		efiApp: efiApp,
	}
}

func Deny(id uuid.UUID, reason string) Message {
	return &denyMessage{
		id:     id,
		reason: reason,
	}
}