story/cmd/sim-device/firmware.go

package main

import (
	"context"
	"crypto/ed25519"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

const maximumSimulatedFirmwareSize = 268435456

type simulatedDeviceCommand struct {
	ID      string         `json:"id"`
	Type    string         `json:"type"`
	Payload map[string]any `json:"payload"`
}

type simulatedFirmwareManifest struct {
	ReleaseID    string
	Version      string
	Digest       string
	Signature    []byte
	SigningKeyID string
	DownloadPath string
}

type firmwareExecutionOutcome string

const (
	firmwareCompleted      firmwareExecutionOutcome = "completed"
	firmwareFailed         firmwareExecutionOutcome = "failed"
	firmwareInterrupted    firmwareExecutionOutcome = "interrupted"
	firmwareRebootRequired firmwareExecutionOutcome = "reboot_required"
)

func (s *simulator) executeFirmwareCommand(ctx context.Context, command simulatedDeviceCommand) (firmwareExecutionOutcome, error) {
	manifest, err := parseSimulatedFirmwareManifest(command.Payload)
	if err != nil {
		if ackErr := s.ackDeviceCommand(ctx, command.ID, "FAILED", map[string]any{"result_code": "MANIFEST_INVALID", "error_message": err.Error()}); ackErr != nil {
			return firmwareFailed, ackErr
		}
		return firmwareFailed, nil
	}
	fault, err := simulatedFirmwareFault()
	if err != nil {
		return firmwareFailed, err
	}
	delay, err := simulatorOptionalDuration("SIM_FIRMWARE_STEP_DELAY", 30*time.Second)
	if err != nil {
		return firmwareFailed, err
	}
	abEnabled, err := simulatedABFirmwareEnabled()
	if err != nil {
		return firmwareFailed, err
	}
	if !abEnabled && isABFirmwareFault(fault) {
		return firmwareFailed, errors.New("A/B firmware fault injection requires SIM_FIRMWARE_AB_UPDATE=true")
	}
	if abEnabled {
		if handled, outcome, resumeErr := s.resumeABFirmwareCommand(ctx, command, manifest, fault, delay); handled || resumeErr != nil {
			return outcome, resumeErr
		}
	}
	if err := s.ackDeviceCommand(ctx, command.ID, "ACKNOWLEDGED", map[string]any{"phase": "downloading", "progress_percent": 5, "result_code": "DOWNLOAD_STARTED"}); err != nil {
		return firmwareFailed, err
	}
	if fault == "download_failed" {
		return firmwareFailed, s.ackDeviceCommand(ctx, command.ID, "FAILED", map[string]any{"phase": "downloading", "progress_percent": 5, "result_code": "DOWNLOAD_FAILED"})
	}
	content, err := s.downloadFirmware(ctx, manifest.DownloadPath)
	if err != nil {
		if ackErr := s.ackDeviceCommand(ctx, command.ID, "FAILED", map[string]any{"phase": "downloading", "progress_percent": 5, "result_code": "DOWNLOAD_FAILED", "error_message": err.Error()}); ackErr != nil {
			return firmwareFailed, ackErr
		}
		return firmwareFailed, nil
	}
	digestValue := sha256.Sum256(content)
	actualDigest := hex.EncodeToString(digestValue[:])
	if fault == "digest_mismatch" || actualDigest != manifest.Digest {
		return firmwareFailed, s.ackDeviceCommand(ctx, command.ID, "FAILED", map[string]any{"phase": "downloading", "progress_percent": 100, "result_code": "DIGEST_MISMATCH", "expected_sha256": manifest.Digest, "actual_sha256": actualDigest, "bytes_downloaded": len(content)})
	}
	publicKey, err := simulatedFirmwareTrustedKey(manifest.SigningKeyID)
	if err != nil {
		return firmwareFailed, s.ackDeviceCommand(ctx, command.ID, "FAILED", map[string]any{"phase": "downloading", "progress_percent": 100, "result_code": "SIGNING_KEY_UNTRUSTED", "signing_key_id": manifest.SigningKeyID, "error_message": err.Error()})
	}
	signedManifest := fmt.Sprintf("firmware-manifest-v1\nmodel=%s\nhardware_revision=%s\nversion=%s\nsha256=%s\nsize=%d\n", env("SIM_DEVICE_MODEL", "simulated-terminal"), env("SIM_DEVICE_HARDWARE_REVISION", "sim-rev-a"), manifest.Version, manifest.Digest, len(content))
	if fault == "signature_mismatch" || !ed25519.Verify(publicKey, []byte(signedManifest), manifest.Signature) {
		return firmwareFailed, s.ackDeviceCommand(ctx, command.ID, "FAILED", map[string]any{"phase": "downloading", "progress_percent": 100, "result_code": "SIGNATURE_MISMATCH", "signing_key_id": manifest.SigningKeyID, "verified_sha256": actualDigest, "bytes_downloaded": len(content)})
	}
	if err := s.ackDeviceCommand(ctx, command.ID, "ACKNOWLEDGED", map[string]any{"phase": "downloading", "progress_percent": 100, "result_code": "ARTIFACT_VERIFIED", "verified_sha256": actualDigest, "verified_signing_key_id": manifest.SigningKeyID, "bytes_downloaded": len(content)}); err != nil {
		return firmwareFailed, err
	}
	if fault == "interrupt_after_download" {
		return firmwareInterrupted, nil
	}
	if err := waitForContext(ctx, delay); err != nil {
		return firmwareInterrupted, err
	}
	if err := s.ackDeviceCommand(ctx, command.ID, "ACKNOWLEDGED", map[string]any{"phase": "installing", "progress_percent": 80, "result_code": "INSTALL_STARTED"}); err != nil {
		return firmwareFailed, err
	}
	if fault == "install_failed" {
		return firmwareFailed, s.ackDeviceCommand(ctx, command.ID, "FAILED", map[string]any{"phase": "installing", "progress_percent": 80, "result_code": "INSTALL_FAILED"})
	}
	if fault == "interrupt_after_install" {
		return firmwareInterrupted, nil
	}
	if err := waitForContext(ctx, delay); err != nil {
		return firmwareInterrupted, err
	}
	if abEnabled {
		return s.stageABFirmwareCommand(ctx, command, manifest, fault)
	}
	if err := s.ackDeviceCommand(ctx, command.ID, "SUCCEEDED", map[string]any{"phase": "confirmed", "progress_percent": 100, "result_code": "SIMULATED_OK", "installed_version": manifest.Version, "verified_sha256": actualDigest, "verified_signing_key_id": manifest.SigningKeyID}); err != nil {
		return firmwareFailed, err
	}
	return firmwareCompleted, nil
}

func parseSimulatedFirmwareManifest(payload map[string]any) (simulatedFirmwareManifest, error) {
	digest := stringValue(payload, "sha256")
	signature, signatureErr := base64.StdEncoding.DecodeString(stringValue(payload, "signature"))
	manifest := simulatedFirmwareManifest{
		ReleaseID:    stringValue(payload, "release_id"),
		Version:      stringValue(payload, "version"),
		Digest:       digest,
		Signature:    signature,
		SigningKeyID: stringValue(payload, "signing_key_id"),
		DownloadPath: stringValue(payload, "download_path"),
	}
	if manifest.ReleaseID == "" || manifest.Version == "" || manifest.SigningKeyID == "" || len(manifest.Digest) != sha256.Size*2 || signatureErr != nil || len(manifest.Signature) != ed25519.SignatureSize {
		return simulatedFirmwareManifest{}, errors.New("firmware command requires release_id, version, signing_key_id, a 64-character SHA-256 and a 64-byte Ed25519 signature")
	}
	if _, err := hex.DecodeString(manifest.Digest); err != nil {
		return simulatedFirmwareManifest{}, errors.New("firmware command SHA-256 is not lowercase hexadecimal")
	}
	if manifest.Digest != strings.ToLower(manifest.Digest) {
		return simulatedFirmwareManifest{}, errors.New("firmware command SHA-256 is not lowercase hexadecimal")
	}
	parsedPath, err := url.Parse(manifest.DownloadPath)
	if err != nil || parsedPath.IsAbs() || parsedPath.Host != "" || parsedPath.RawQuery != "" || parsedPath.Fragment != "" || strings.Contains(parsedPath.Path, "..") || !strings.HasPrefix(parsedPath.Path, "/device/v2/firmware-releases/") || !strings.HasSuffix(parsedPath.Path, "/download") {
		return simulatedFirmwareManifest{}, errors.New("firmware command download_path is outside the fixed device firmware boundary")
	}
	return manifest, nil
}

func stringValue(value map[string]any, key string) string {
	item, _ := value[key].(string)
	return strings.TrimSpace(item)
}

func simulatedFirmwareFault() (string, error) {
	value := strings.ToLower(strings.TrimSpace(os.Getenv("SIM_FIRMWARE_FAULT")))
	if value == "" || value == "none" {
		return "none", nil
	}
	switch value {
	case "download_failed", "digest_mismatch", "signature_mismatch", "install_failed", "interrupt_after_download", "interrupt_after_install", "power_loss_after_stage", "power_loss_during_boot", "boot_health_failed":
		return value, nil
	default:
		return "", fmt.Errorf("SIM_FIRMWARE_FAULT must be none, download_failed, digest_mismatch, signature_mismatch, install_failed, interrupt_after_download, interrupt_after_install, power_loss_after_stage, power_loss_during_boot or boot_health_failed")
	}
}

func isABFirmwareFault(value string) bool {
	return value == "power_loss_after_stage" || value == "power_loss_during_boot" || value == "boot_health_failed"
}

func simulatedFirmwareTrustedKey(keyID string) (ed25519.PublicKey, error) {
	raw := strings.TrimSpace(os.Getenv("SIM_FIRMWARE_TRUSTED_KEYS"))
	if raw == "" {
		return nil, errors.New("SIM_FIRMWARE_TRUSTED_KEYS is empty")
	}
	var encoded map[string]string
	if err := json.Unmarshal([]byte(raw), &encoded); err != nil {
		return nil, errors.New("SIM_FIRMWARE_TRUSTED_KEYS must be a JSON object of key IDs to Base64 Ed25519 public keys")
	}
	value, exists := encoded[keyID]
	if !exists {
		return nil, fmt.Errorf("signing key %s is not provisioned", keyID)
	}
	publicKey, err := base64.StdEncoding.DecodeString(value)
	if err != nil || len(publicKey) != ed25519.PublicKeySize {
		return nil, fmt.Errorf("signing key %s is not a 32-byte Base64 Ed25519 public key", keyID)
	}
	return ed25519.PublicKey(publicKey), nil
}

func simulatorOptionalDuration(key string, maximum time.Duration) (time.Duration, error) {
	value := strings.TrimSpace(os.Getenv(key))
	if value == "" {
		return 0, nil
	}
	parsed, err := time.ParseDuration(value)
	if err != nil || parsed < 0 || parsed > maximum {
		return 0, fmt.Errorf("%s must be a non-negative duration no greater than %s", key, maximum)
	}
	return parsed, nil
}

func (s *simulator) ackDeviceCommand(ctx context.Context, commandID, status string, result map[string]any) error {
	body, err := json.Marshal(map[string]any{"status": status, "result_summary": result})
	if err != nil {
		return err
	}
	_, err = s.request(ctx, http.MethodPost, "/device/v2/commands/"+commandID+"/ack", body, "")
	return err
}

func (s *simulator) downloadFirmware(ctx context.Context, path string) ([]byte, error) {
	result, err := s.request(ctx, http.MethodGet, path, nil, "")
	if err != nil {
		return nil, err
	}
	var address struct {
		DownloadURL string `json:"download_url"`
	}
	if err := decodeData(result.Body, &address); err != nil || address.DownloadURL == "" {
		return nil, errors.New("firmware download response did not contain a private URL")
	}
	parsed, err := url.Parse(address.DownloadURL)
	if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.Host == "" {
		return nil, errors.New("firmware private download URL is invalid")
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, address.DownloadURL, nil)
	if err != nil {
		return nil, err
	}
	response, err := s.client.Do(request)
	if err != nil {
		return nil, fmt.Errorf("download firmware object: %w", err)
	}
	defer response.Body.Close()
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		return nil, fmt.Errorf("firmware object returned HTTP %d", response.StatusCode)
	}
	content, err := io.ReadAll(io.LimitReader(response.Body, maximumSimulatedFirmwareSize+1))
	if err != nil {
		return nil, fmt.Errorf("read firmware object: %w", err)
	}
	if len(content) == 0 || len(content) > maximumSimulatedFirmwareSize {
		return nil, fmt.Errorf("firmware object size %d is outside 1 to %d bytes", len(content), maximumSimulatedFirmwareSize)
	}
	return content, nil
}