a73x

internal/cloudinit/multipart.go

Ref:   Size: 11.6 KiB   History

package cloudinit

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"net/mail"
	"net/textproto"
	"regexp"
	"strings"
)

// format classifies cloud-init user-data. cloud-init picks a handler by the
// payload's leading "magic" line (or MIME/gzip framing), so user-data is a
// tagged union, not just YAML — detection mirrors cloud-init's own sniff.
type format int

const (
	formatUnknown format = iota
	formatCloudConfig
	formatShellScript
	formatBoothook
	formatInclude
	formatPartHandler
	formatMultipart // MIME multipart/mixed archive
	formatJinja     // ## template: jinja — templated; not safe to edit or wrap blind
	formatGzip      // gzip-compressed payload
)

func (f format) String() string {
	switch f {
	case formatCloudConfig:
		return "cloud-config"
	case formatShellScript:
		return "shell-script"
	case formatBoothook:
		return "cloud-boothook"
	case formatInclude:
		return "include"
	case formatPartHandler:
		return "part-handler"
	case formatMultipart:
		return "multipart"
	case formatJinja:
		return "jinja-template"
	case formatGzip:
		return "gzip"
	default:
		return "unknown"
	}
}

// headerRe matches an RFC 5322 header field name at the start of a line —
// used to tell a MIME archive (starts with headers) from a #-tagged payload.
var headerRe = regexp.MustCompile(`^[A-Za-z][A-Za-z0-9-]*:`)

// detectFormat classifies user-data the way cloud-init does: gzip magic first,
// then the leading non-blank line's marker (its first whitespace-delimited
// token, so a trailing comment is tolerated), then — for a payload that leads
// with RFC 5322 headers rather than a #-marker — a real MIME header-block parse
// (an archive may put MIME-Version before Content-Type).
func detectFormat(userData string) format {
	if len(userData) >= 2 && userData[0] == 0x1f && userData[1] == 0x8b {
		return formatGzip
	}
	first := firstNonBlankLine(userData)
	if strings.HasPrefix(first, "## template: jinja") {
		return formatJinja
	}
	switch marker := firstMarker(userData); {
	case marker == "#cloud-config":
		return formatCloudConfig
	case strings.HasPrefix(marker, "#!"):
		return formatShellScript
	case marker == "#cloud-boothook":
		return formatBoothook
	case strings.HasPrefix(marker, "#include"): // #include and #include-once
		return formatInclude
	case marker == "#part-handler":
		return formatPartHandler
	}
	if isMIMEMultipart(userData) {
		return formatMultipart
	}
	return formatUnknown
}

// isMIMEMultipart reports whether s is a MIME multipart archive by parsing its
// header block. It only tries when the first line looks like a header and is
// NOT a #-tagged payload — otherwise a cloud-config line such as `packages:`
// would masquerade as an RFC 5322 header.
func isMIMEMultipart(s string) bool {
	first := firstNonBlankLine(s)
	if strings.HasPrefix(first, "#") || !headerRe.MatchString(first) {
		return false
	}
	msg, err := mail.ReadMessage(strings.NewReader(s))
	if err != nil {
		return false
	}
	mt, _, err := mime.ParseMediaType(msg.Header.Get("Content-Type"))
	return err == nil && strings.HasPrefix(mt, "multipart/")
}

// AddSSHKey returns user-data with key installed for the default user,
// regardless of the user-data's format, WITHOUT editing formats where an
// in-place edit is unsafe:
//   - cloud-config → merged into ssh_authorized_keys (a clean single document).
//   - script / boothook / include / part-handler → wrapped in a multipart
//     archive alongside a small #cloud-config key part; the original payload is
//     left byte-for-byte intact as its own part.
//   - multipart → the key part is APPENDED to the existing archive (not nested),
//     preserving every original part's body AND headers.
//   - jinja / gzip / unknown → error: we will not blindly edit a templated or
//     opaque payload, so the caller surfaces a clear failure instead of a
//     silent no-op.
//
// key must already be validated single-line (callers do): it is embedded in a
// generated #cloud-config, so a newline would produce a garbage key. It is
// added as a YAML scalar node, so it cannot inject structure regardless.
func AddSSHKey(userData, key string) (string, error) {
	f := detectFormat(userData)
	if f == formatCloudConfig {
		// The common case is a single clean document, so it stays one: merging
		// the key in beats handing cloud-init an archive to reassemble.
		return mergeSSHKey(userData, key)
	}
	out, err := appendPart(userData, f, keyPart(key))
	if err != nil {
		return "", fmt.Errorf("cannot add an ssh key to %s user-data; include the key in the user-data itself", f)
	}
	return out, nil
}

// AppendCloudConfig returns userData with doc — a #cloud-config document eitri
// generated — appended as the LAST part of a MIME multipart archive, leaving
// whatever the tenant supplied byte-for-byte intact as the earlier part(s).
//
// Last is the contract, not a detail. cloud-init's cloud-config part handler
// merges each part into the buffer the earlier parts built, using the mergers
// named by THAT part (its own `merge_how`, or a `Merge-Type` MIME header), and
// falling back to `dict(replace)+list()+str()` for a part that names none. So a
// document that carries a merge_how governs only what is merged into it — it
// has to arrive after the document it means to merge with. This is the layer
// where merge_how is read at all: the cloud-config handler POPS the key out of
// each part as it processes it and writes the stripped result, so the later
// user-data-versus-vendor-data merge never sees one (see mergeHow in
// internal/agent/seed).
//
// A tenant document is never edited to make room. A plain cloud-config becomes
// its own part; a script, boothook, include or part-handler payload rides under
// the type cloud-init dispatches it with; an archive the tenant already built
// gets the part appended to it rather than nested inside a new one, keeping
// every original part's body AND headers.
//
// Jinja, gzip and unrecognised payloads are refused. Their real content-type is
// not knowable without rendering or decompressing them, and a part labelled
// wrongly is a part cloud-init runs wrongly — where the alternative is handing
// a guest a seed whose CA trust silently never lands, an error at build time is
// the only honest answer.
func AppendCloudConfig(userData, doc string) (string, error) {
	f := detectFormat(userData)
	out, err := appendPart(userData, f, typedPart("text/cloud-config", doc))
	if err != nil {
		return "", fmt.Errorf("cannot compose eitri's cloud-config with %s user-data: %w", f, err)
	}
	return out, nil
}

// appendPart returns userData with last appended as the final part of a
// multipart archive: an archive gains a part, anything else is wrapped with
// itself first. f is userData's already-detected format, so callers that
// special-case one of them (AddSSHKey) classify only once.
func appendPart(userData string, f format, last part) (string, error) {
	if f == formatMultipart {
		return appendToMultipart(userData, last)
	}
	ct, ok := partType(f)
	if !ok {
		return "", fmt.Errorf("%s payloads cannot be given a MIME part header", f)
	}
	return wrapMultipart([]part{typedPart(ct, userData), last})
}

// partType maps a detected format to the content-type cloud-init dispatches
// that payload with. The bool reports whether the payload can ride in an
// archive at all: jinja renders to a type nobody knows yet, gzip is opaque
// bytes, and an unrecognised payload names no handler to run it.
func partType(f format) (string, bool) {
	switch f {
	case formatCloudConfig:
		return "text/cloud-config", true
	case formatShellScript:
		return "text/x-shellscript", true
	case formatBoothook:
		return "text/cloud-boothook", true
	case formatInclude:
		return "text/x-include-url", true
	case formatPartHandler:
		return "text/part-handler", true
	}
	return "", false
}

// part is one MIME sub-document: its headers and its (already-decoded) body.
type part struct {
	header textproto.MIMEHeader
	body   string
}

// typedPart builds a part with a single Content-Type (+ charset) and
// MIME-Version — used for eitri-generated parts and simple wraps.
func typedPart(contentType, body string) part {
	h := textproto.MIMEHeader{}
	h.Set("Content-Type", contentType+`; charset="utf-8"`)
	h.Set("MIME-Version", "1.0")
	return part{h, body}
}

// keyPart is the generated #cloud-config carrying just the ssh key, built via
// mergeSSHKey so the key lands in a valid ssh_authorized_keys list.
func keyPart(key string) part {
	doc, _ := mergeSSHKey("#cloud-config\n", key) // cannot fail on a literal #cloud-config
	return typedPart("text/cloud-config", doc)
}

// wrapMultipart serializes parts into a cloud-init multipart/mixed archive:
// a top-level Content-Type/MIME-Version header block, then each part with its
// own headers. cloud-init runs every part by its type and merges the
// cloud-config ones, so eitri's key part composes with the user's payload
// without eitri ever editing that payload.
func wrapMultipart(parts []part) (string, error) {
	var body bytes.Buffer
	mw := multipart.NewWriter(&body)
	for _, p := range parts {
		h := p.header
		if h == nil {
			h = textproto.MIMEHeader{}
		}
		if h.Get("Content-Type") == "" {
			h.Set("Content-Type", "text/plain")
		}
		if h.Get("MIME-Version") == "" {
			h.Set("MIME-Version", "1.0")
		}
		pw, err := mw.CreatePart(h)
		if err != nil {
			return "", fmt.Errorf("multipart part: %w", err)
		}
		if _, err := pw.Write([]byte(p.body)); err != nil {
			return "", fmt.Errorf("multipart write: %w", err)
		}
	}
	if err := mw.Close(); err != nil {
		return "", fmt.Errorf("multipart close: %w", err)
	}
	return "Content-Type: multipart/mixed; boundary=\"" + mw.Boundary() + "\"\n" +
		"MIME-Version: 1.0\n\n" + body.String(), nil
}

// appendToMultipart parses an existing multipart archive, keeps every original
// part's body AND headers (cloud-init uses part headers like
// Content-Disposition/filename to name and order scripts, and Merge-Type to
// control cloud-config merging), and re-emits with last appended —
// deliberately NOT nesting the user's archive inside a new one.
func appendToMultipart(userData string, last part) (string, error) {
	msg, err := mail.ReadMessage(strings.NewReader(userData))
	if err != nil {
		return "", fmt.Errorf("parse multipart headers: %w", err)
	}
	mediaType, params, err := mime.ParseMediaType(msg.Header.Get("Content-Type"))
	if err != nil || !strings.HasPrefix(mediaType, "multipart/") {
		return "", fmt.Errorf("not a multipart archive: %v", err)
	}
	boundary, ok := params["boundary"]
	if !ok {
		return "", errors.New("multipart archive missing boundary")
	}
	mr := multipart.NewReader(msg.Body, boundary)
	var parts []part
	for {
		p, err := mr.NextPart()
		if errors.Is(err, io.EOF) {
			break // normal end of archive
		}
		if err != nil {
			return "", fmt.Errorf("read multipart part: %w", err)
		}
		buf, err := io.ReadAll(p)
		if err != nil {
			return "", fmt.Errorf("read multipart part body: %w", err)
		}
		h := textproto.MIMEHeader{}
		for k, vs := range p.Header {
			// multipart.Reader has already decoded any transfer encoding, so
			// re-emitting Content-Transfer-Encoding would mislabel the decoded
			// body; every other header (Content-Disposition, Merge-Type, …) is
			// preserved.
			if strings.EqualFold(k, "Content-Transfer-Encoding") {
				continue
			}
			for _, v := range vs {
				h.Add(k, v)
			}
		}
		parts = append(parts, part{header: h, body: string(buf)})
	}
	if len(parts) == 0 {
		return "", errors.New("multipart archive has no parts")
	}
	return wrapMultipart(append(parts, last))
}

// firstNonBlankLine returns the first line with non-whitespace content, trimmed.
func firstNonBlankLine(s string) string {
	for line := range strings.SplitSeq(s, "\n") {
		if t := strings.TrimSpace(line); t != "" {
			return t
		}
	}
	return ""
}