internal/cloudinit/cloudinit.go

Ref: Size: 5.8 KiB History
// Package cloudinit merges eitri's structured VM inputs into user-supplied
// cloud-init user-data. It is a neutral leaf: the control plane (server/api)
// uses it at create time so a bad merge fails fast as a 400, and it holds no
// dependency on either plane.
//
// Today it does one thing — install an SSH public key into whatever user-data
// the caller supplied — but it exists as a package because "eitri owns a bit of
// your cloud-init" is a real responsibility that deserves one home and real
// tests, not a string hack buried in a handler.
package cloudinit

import (
	"errors"
	"fmt"
	"io"
	"strings"

	"gopkg.in/yaml.v3"
)

// errNotCloudConfig reports user-data that is not a `#cloud-config` document —
// a shell script, a jinja-templated config, MIME multipart, etc. There is no
// ssh_authorized_keys to merge into, so the caller must reject rather than
// silently drop the key (the exact bug this package removes).
var errNotCloudConfig = errors.New("user-data is not #cloud-config; cannot merge an ssh key into it")

// mergeSSHKey returns userData with key added to the top-level
// ssh_authorized_keys list — cloud-init's canonical way to add a key to the
// image's default user, robust whether or not the doc defines a `users:` block.
// (An operator who REPLACES the default user with their own named user should
// put the key in that users block themselves.)
//
// The edit is done IN PLACE on the YAML node tree, so every other value,
// comment, and key order in the user's document is preserved byte-for-byte —
// only the ssh_authorized_keys sequence is touched (created, or appended to
// with de-duplication; a scalar value is normalized to a sequence first). The
// `#cloud-config` header is re-emitted.
//
// Multi-document YAML (a second `---`) is rejected: cloud-init rejects it too,
// and accepting-then-silently-dropping the later documents would mask an error
// the guest would otherwise raise.
//
// key must already be validated single-line by the caller. Even so, the key is
// added as a YAML scalar node (not string-interpolated), so a value with
// YAML-significant characters — or a bypassed multi-line value — is emitted as
// a quoted/block scalar and cannot inject structure.
func mergeSSHKey(userData, key string) (string, error) {
	if !isCloudConfig(userData) {
		return "", errNotCloudConfig
	}

	dec := yaml.NewDecoder(strings.NewReader(userData))
	var doc yaml.Node
	switch err := dec.Decode(&doc); {
	case errors.Is(err, io.EOF):
		// A #cloud-config with only comments / no body: synthesize an empty
		// mapping document to hang the key on.
		doc = yaml.Node{Kind: yaml.DocumentNode, Content: []*yaml.Node{{Kind: yaml.MappingNode}}}
	case err != nil:
		return "", fmt.Errorf("parse cloud-config: %w", err)
	default:
		// Reject a second document (cloud-init rejects multi-doc cloud-config).
		var extra yaml.Node
		if err := dec.Decode(&extra); err == nil {
			return "", errors.New("multi-document cloud-config is not supported")
		} else if !errors.Is(err, io.EOF) {
			return "", fmt.Errorf("parse cloud-config: %w", err)
		}
	}

	if len(doc.Content) == 0 {
		doc.Content = []*yaml.Node{{Kind: yaml.MappingNode}}
	}
	root := doc.Content[0]
	if root.Kind != yaml.MappingNode {
		return "", fmt.Errorf("cloud-config root must be a mapping, not %s", kindName(root.Kind))
	}
	if err := appendAuthorizedKey(root, key); err != nil {
		return "", err
	}

	out, err := yaml.Marshal(&doc)
	if err != nil {
		return "", fmt.Errorf("marshal cloud-config: %w", err)
	}
	return "#cloud-config\n" + string(out), nil
}

// appendAuthorizedKey adds key to the mapping's ssh_authorized_keys sequence,
// creating the entry if absent, normalizing a scalar to a sequence, and
// de-duplicating. Every other node in the mapping is left untouched.
func appendAuthorizedKey(root *yaml.Node, key string) error {
	keyNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: key}
	for i := 0; i+1 < len(root.Content); i += 2 {
		if root.Content[i].Value != "ssh_authorized_keys" {
			continue
		}
		val := root.Content[i+1]
		switch val.Kind {
		case yaml.SequenceNode:
			for _, e := range val.Content {
				if e.Value == key {
					return nil // already present
				}
			}
			val.Content = append(val.Content, keyNode)
		case yaml.ScalarNode:
			if val.Value == key {
				return nil
			}
			existing := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: val.Value}
			*val = yaml.Node{Kind: yaml.SequenceNode, Content: []*yaml.Node{existing, keyNode}}
		default:
			return fmt.Errorf("ssh_authorized_keys must be a string or list, not %s", kindName(val.Kind))
		}
		return nil
	}
	// Absent: append `ssh_authorized_keys: [key]`.
	root.Content = append(root.Content,
		&yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: "ssh_authorized_keys"},
		&yaml.Node{Kind: yaml.SequenceNode, Content: []*yaml.Node{keyNode}})
	return nil
}

// isCloudConfig reports whether s is a cloud-config document: its first
// non-blank line's first whitespace-delimited token must be exactly
// `#cloud-config` (so a trailing comment like `#cloud-config # my vm` still
// counts, but `#cloud-config-archive` — a different format — does not). A
// `## template: jinja` preamble or any other first line is rejected.
func isCloudConfig(s string) bool {
	return firstMarker(s) == "#cloud-config"
}

// firstMarker returns the first whitespace-delimited token of the first
// non-blank line — cloud-init's per-format "magic" marker.
func firstMarker(s string) string {
	if f := strings.Fields(firstNonBlankLine(s)); len(f) > 0 {
		return f[0]
	}
	return ""
}

func kindName(k yaml.Kind) string {
	switch k {
	case yaml.DocumentNode:
		return "document"
	case yaml.SequenceNode:
		return "list"
	case yaml.MappingNode:
		return "mapping"
	case yaml.ScalarNode:
		return "scalar"
	case yaml.AliasNode:
		return "alias"
	default:
		return "unknown"
	}
}