internal/cloudinit/multipart_test.go
Ref: Size: 6.4 KiB History
package cloudinit
import (
"io"
"mime"
"mime/multipart"
"net/mail"
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestDetectFormat(t *testing.T) {
cases := []struct {
in string
want format
}{
{"#cloud-config\npackages: [htop]\n", formatCloudConfig},
{"\n\n#cloud-config\n", formatCloudConfig}, // leading blank lines tolerated
{"#!/bin/bash\necho hi\n", formatShellScript},
{"#cloud-boothook\n#!/bin/sh\n", formatBoothook},
{"#include\nhttps://example/x\n", formatInclude},
{"#include-once\nhttps://example/x\n", formatInclude},
{"#part-handler\n", formatPartHandler},
{"## template: jinja\n#cloud-config\n", formatJinja},
{"Content-Type: multipart/mixed; boundary=\"X\"\n\n", formatMultipart},
{"\x1f\x8b\x08 gzip bytes", formatGzip},
{"just some text", formatUnknown},
}
for _, c := range cases {
assert.Equal(t, c.want, detectFormat(c.in), "input %q", c.in)
}
}
// mimeParts parses a cloud-init multipart archive back into (contentType, body)
// pairs so tests assert on STRUCTURE, not on the random MIME boundary.
func mimeParts(t *testing.T, s string) map[string]string {
t.Helper()
msg, err := mail.ReadMessage(strings.NewReader(s))
require.NoError(t, err)
mediaType, params, err := mime.ParseMediaType(msg.Header.Get("Content-Type"))
require.NoError(t, err)
require.Equal(t, "multipart/mixed", mediaType)
mr := multipart.NewReader(msg.Body, params["boundary"])
out := map[string]string{}
for {
p, err := mr.NextPart()
if err != nil {
break
}
body, err := io.ReadAll(p)
require.NoError(t, err)
mt, _, _ := mime.ParseMediaType(p.Header.Get("Content-Type"))
out[mt] = string(body)
}
return out
}
func TestAddSSHKeyCloudConfigMerges(t *testing.T) {
// cloud-config path delegates to mergeSSHKey → a single merged document,
// NOT a multipart wrapper.
out, err := AddSSHKey("#cloud-config\npackages:\n - htop\n", key)
require.NoError(t, err)
assert.True(t, strings.HasPrefix(out, "#cloud-config\n"), "stays a single cloud-config doc, got: %q", out[:min(24, len(out))])
assert.Contains(t, out, key)
assert.NotContains(t, out, "multipart", "the common case must not become a MIME archive")
}
func TestAddSSHKeyShellScriptWraps(t *testing.T) {
script := "#!/bin/bash\necho hello > /tmp/marker\n"
out, err := AddSSHKey(script, key)
require.NoError(t, err)
parts := mimeParts(t, out)
// The script survives byte-for-byte as its own part...
assert.Equal(t, script, parts["text/x-shellscript"], "the user's script must be intact and untouched")
// ...and the key rides a sibling cloud-config part.
require.Contains(t, parts, "text/cloud-config")
assert.Contains(t, parts["text/cloud-config"], "#cloud-config")
assert.Contains(t, parts["text/cloud-config"], key)
}
func TestAddSSHKeyBoothookWraps(t *testing.T) {
bh := "#cloud-boothook\n#!/bin/sh\necho early\n"
out, err := AddSSHKey(bh, key)
require.NoError(t, err)
parts := mimeParts(t, out)
assert.Equal(t, bh, parts["text/cloud-boothook"])
assert.Contains(t, parts["text/cloud-config"], key)
}
func TestAddSSHKeyAppendsToExistingMultipart_NotNested(t *testing.T) {
// A user who already supplied a multipart archive: the key part is appended
// alongside their parts, and their parts survive verbatim — no nesting.
existing, err := wrapMultipart([]part{
typedPart("text/cloud-config", "#cloud-config\npackages:\n - git\n"),
typedPart("text/x-shellscript", "#!/bin/bash\necho hi\n"),
})
require.NoError(t, err)
out, err := AddSSHKey(existing, key)
require.NoError(t, err)
parts := mimeParts(t, out)
// No part is itself a multipart archive → not nested.
for ct := range parts {
assert.NotContains(t, ct, "multipart", "must append, not nest")
}
assert.Contains(t, parts["text/x-shellscript"], "echo hi", "original script part survives")
assert.Contains(t, parts["text/cloud-config"], key, "key part added")
}
func TestAddSSHKeyRejectsJinjaAndGzipAndUnknown(t *testing.T) {
for _, in := range []string{
"## template: jinja\n#cloud-config\nhostname: {{ v1.local_hostname }}\n",
"\x1f\x8b\x08 gzipped",
"random text that is not user-data",
} {
_, err := AddSSHKey(in, key)
assert.Error(t, err, "must reject %q rather than silently no-op", in[:min(20, len(in))])
}
}
func TestDetectFormatMultipartHeaderOrder(t *testing.T) {
// A valid archive may lead with MIME-Version before Content-Type; detection
// must parse the header block, not just sniff line one.
in := "MIME-Version: 1.0\nContent-Type: multipart/mixed; boundary=\"X\"\n\n--X--\n"
assert.Equal(t, formatMultipart, detectFormat(in))
}
func TestDetectFormatCloudConfigTrailingComment(t *testing.T) {
// cloud-init detects #cloud-config by prefix, so a trailing comment on the
// marker line is valid — don't reject it.
assert.Equal(t, formatCloudConfig, detectFormat("#cloud-config # my vm\npackages: [htop]\n"))
}
func TestAppendPreservesPartHeaders(t *testing.T) {
// A user's multipart part may carry Content-Disposition (filename) that
// cloud-init uses to name/order scripts — appendToMultipart must not drop it.
existing := "Content-Type: multipart/mixed; boundary=\"BOUND\"\nMIME-Version: 1.0\n\n" +
"--BOUND\n" +
"Content-Type: text/x-shellscript; charset=\"utf-8\"\n" +
"Content-Disposition: attachment; filename=\"setup.sh\"\n" +
"MIME-Version: 1.0\n\n" +
"#!/bin/bash\necho hi\n" +
"--BOUND--\n"
out, err := AddSSHKey(existing, key)
require.NoError(t, err)
assert.Contains(t, out, `filename="setup.sh"`, "Content-Disposition/filename must survive the append")
assert.Contains(t, out, key)
}
func TestAppendRejectsMalformedMidStream(t *testing.T) {
// A genuinely broken boundary mid-archive should error, not silently
// truncate to the parts read so far.
in := "Content-Type: multipart/mixed; boundary=\"BOUND\"\nMIME-Version: 1.0\n\n" +
"--BOUND\nContent-Type: text/plain\n\nbody without a closing boundary\n"
_, err := AddSSHKey(in, key)
assert.Error(t, err)
}
func TestWrapMultipartIsWellFormed(t *testing.T) {
// The emitted archive must parse with the stdlib MIME reader (well-formed
// headers, boundary, MIME-Version) — a proxy for cloud-init accepting it.
out, err := wrapMultipart([]part{
typedPart("text/x-shellscript", "#!/bin/bash\n"),
typedPart("text/cloud-config", "#cloud-config\n"),
})
require.NoError(t, err)
assert.True(t, strings.HasPrefix(out, "Content-Type: multipart/mixed; boundary="))
assert.Contains(t, out, "MIME-Version: 1.0")
parts := mimeParts(t, out)
assert.Len(t, parts, 2)
}