From 76f99303deeba6624b7fddcafeb04e4680d7d2e2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Clerget?= Date: Wed, 21 Oct 2026 12:04:06 -0500 Subject: [PATCH] Verify the multi-host staging directory and fail closed on digest mismatch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The launcher staged the WALDO binary, worker state, and resume checkpoints beneath the configured scratch path on every secondary host without checking who could modify that path. Another account that could write to the directory or any ancestor, for example by creating /tmp/waldo- first, could replace the binary before it ran with the launching user's identity. The staging commands were also chained with ';', so a SHA-256 mismatch never stopped the rename that published the binary or checkpoint. Before staging, verify that the session directory, every ancestor, and every ancestor of its resolved path are owned by the remote user or root. Directories owned by the remote user must not be writable by their group or other users; on Linux the group bits carry the POSIX ACL mask, so ACL entries cannot grant write either. Root-owned directories writable by others must be sticky. Uploads go to a unique mktemp file within the verified directory and the script exits before the rename on a digest mismatch. An already staged checkpoint is trusted only inside a verified directory. The session directory, cache arguments, delayed worker startup, worker PID handling, and cleanup are unchanged. The checks use only portable tools (ls -ldn, cd -P, mktemp, sha256sum or shasum), so the staging tests run the real scripts through a local ssh shim on Linux and macOS. Fixes #18 Signed-off-by: Cédric Clerget --- docs/MULTI-NODE-TRAINING.md | 17 +++ internal/cli/model_hostfile.go | 103 ++++++++++++-- internal/cli/model_hostfile_test.go | 206 +++++++++++++++++++++++++++- 3 files changed, 313 insertions(+), 13 deletions(-) diff --git a/docs/MULTI-NODE-TRAINING.md b/docs/MULTI-NODE-TRAINING.md index 75c1a359..1aad3a2b 100644 --- a/docs/MULTI-NODE-TRAINING.md +++ b/docs/MULTI-NODE-TRAINING.md @@ -62,6 +62,23 @@ launch. The binary, worker scratch, and resume staging are removed when the session finishes or aborts. Configure the same absolute scratch path on storage with enough capacity on every node; launcher state does not use `/tmp`. +Uploads go to a unique temporary file in the session directory and are +published only after their SHA-256 matches; a mismatch leaves nothing behind. +WALDO refuses to stage the binary, worker state, or checkpoints unless the +session directory, every ancestor, and every ancestor of its resolved path are +owned by the remote user or root and cannot be modified by another account: + +- Directories owned by the remote user must not be writable by their group or + other users. Remove such access with `chmod go-w`. On Linux the group + permission bits include the POSIX ACL mask, so named ACL entries cannot grant + write access either. +- Root-owned directories writable by other users, such as `/tmp`, must be + sticky. + +Staging fails and names the offending path otherwise. If another account +created a directory on the scratch path first, an administrator must remove it +before retrying. + WALDO does not install or modify GPU drivers, CUDA, NCCL, Python, PyTorch, or TorchTitan. Those machine-level runtimes must be installed before training. diff --git a/internal/cli/model_hostfile.go b/internal/cli/model_hostfile.go index 382d8933..a2036629 100644 --- a/internal/cli/model_hostfile.go +++ b/internal/cli/model_hostfile.go @@ -470,16 +470,103 @@ func (session *hostfileSession) remoteCommandContext(ctx context.Context, argume return command } +// remoteStagingFunctions defines POSIX shell functions for staging files on +// secondary hosts. waldo_private_dir creates an absolute directory and +// verifies that it, every ancestor, and every ancestor of its resolved path +// are owned by the remote user or root and cannot be modified by another +// account. Directories owned by the remote user must not be writable by their +// group or other users. Root-owned directories writable by others, such as +// /tmp, must be sticky. On Linux the group permission bits report the POSIX ACL +// mask, which bounds every named ACL entry; other systems' extended ACLs cannot +// be evaluated from the mode and are refused. Only portable tools are used so +// the same script runs on Linux hosts and in macOS unit tests. +const remoteStagingFunctions = `waldo_uid=$(id -u) +waldo_system=$(uname -s) +waldo_check_entry() { + waldo_listing=$(LC_ALL=C ls -ldn -- "$1") || return 1 + set -- "$1" $waldo_listing + case $4 in + "$waldo_uid") waldo_owner=user ;; + 0) waldo_owner=root ;; + *) echo "$1 is owned by uid $4, not uid $waldo_uid or root" >&2; return 1 ;; + esac + case $2 in l*) return 0 ;; esac + case $2 in + ??????????+*) + if [ "$waldo_system" != Linux ]; then + echo "$1 has an access control list that cannot be verified" >&2 + return 1 + fi + ;; + esac + case $2 in ?????w*|????????w*) ;; *) return 0 ;; esac + if [ "$waldo_owner" = root ]; then + case $2 in ?????????[tT]*) return 0 ;; esac + echo "$1 is writable by other users and is not sticky" >&2 + return 1 + fi + echo "$1 is writable by its group or other users; remove that access with chmod go-w" >&2 + return 1 +} +waldo_check_ancestors() { + waldo_path=$1 + while :; do + waldo_check_entry "$waldo_path" || return 1 + [ "$waldo_path" = / ] && return 0 + waldo_path=$(dirname -- "$waldo_path") + done +} +waldo_private_dir() { + case $1 in /*) ;; *) echo "staging directory $1 is not absolute" >&2; return 1 ;; esac + mkdir -p -- "$1" || return 1 + waldo_physical=$(cd -P -- "$1" && pwd -P) || return 1 + waldo_check_ancestors "$1" || return 1 + waldo_check_ancestors "$waldo_physical" +} +waldo_sha256() { + if command -v sha256sum >/dev/null 2>&1; then + waldo_sum=$(sha256sum -- "$1") + else + waldo_sum=$(shasum -a 256 -- "$1") + fi + printf '%s\n' "${waldo_sum%% *}" +} +` + +// remoteStageScript returns a script that receives a file on standard input +// and publishes it at target only after its SHA-256 matches digest. The upload +// goes to a unique temporary file inside the verified private directory and +// is removed on any failure. +func remoteStageScript(target, digest, mode, label string) string { + directory := shellQuote(filepath.Dir(target)) + sum := shellQuote(digest) + return "set -euf\numask 077\n" + remoteStagingFunctions + + "waldo_private_dir " + directory + "\n" + + "tmp=$(mktemp " + directory + "/.waldo-transfer.XXXXXX)\n" + + "trap 'rm -f -- \"$tmp\"' EXIT\n" + + "cat > \"$tmp\"\n" + + "chmod " + mode + " \"$tmp\"\n" + + "sum=$(waldo_sha256 \"$tmp\")\n" + + "if [ \"$sum\" != " + sum + " ]; then echo \"staged " + label + " SHA-256 $sum does not match \"" + sum + " >&2; exit 1; fi\n" + + "mv -f -- \"$tmp\" " + shellQuote(target) + "\n" +} + +// remoteVerifiedFileScript succeeds only when target already exists inside a +// verified private directory with the expected SHA-256. +func remoteVerifiedFileScript(target, digest string) string { + return "set -euf\numask 077\n" + remoteStagingFunctions + + "waldo_private_dir " + shellQuote(filepath.Dir(target)) + "\n" + + "test -f " + shellQuote(target) + "\n" + + "test \"$(waldo_sha256 " + shellQuote(target) + ")\" = " + shellQuote(digest) + "\n" +} + func (session *hostfileSession) stageBinary(host string) error { file, err := os.Open(session.binary) if err != nil { return err } defer file.Close() - temporary := session.remoteBinary + ".tmp" - remote := fmt.Sprintf("umask 077; mkdir -p %s; cat > %s; chmod 700 %s; test \"$(sha256sum %s | cut -d' ' -f1)\" = %s; mv -f %s %s", - shellQuote(session.remoteRoot), shellQuote(temporary), shellQuote(temporary), shellQuote(temporary), shellQuote(session.binarySHA256), shellQuote(temporary), shellQuote(session.remoteBinary)) - command := session.remoteCommand(host, remote) + command := session.remoteCommand(host, remoteStageScript(session.remoteBinary, session.binarySHA256, "700", "WALDO")) command.Stdin = file var output strings.Builder command.Stdout, command.Stderr = &output, &output @@ -814,18 +901,14 @@ func (session *hostfileSession) stageResume(host string, artifacts []training.Ar for index, artifact := range artifacts { localPath := sources[index] remotePath := targets[index] - check := fmt.Sprintf("test -f %s; test \"$(sha256sum %s | cut -d' ' -f1)\" = %s", shellQuote(remotePath), shellQuote(remotePath), shellQuote(artifact.SHA256)) - if err := session.remoteCommand(host, check).Run(); err == nil { + if err := session.remoteCommand(host, remoteVerifiedFileScript(remotePath, artifact.SHA256)).Run(); err == nil { continue } file, err := os.Open(localPath) if err != nil { return fmt.Errorf("stage checkpoint on %s: open %s: %w", host, artifact.Path, err) } - temporary := remotePath + ".waldo-transfer" - remote := fmt.Sprintf("umask 077; mkdir -p %s; cat > %s; test \"$(sha256sum %s | cut -d' ' -f1)\" = %s; mv -f %s %s", - shellQuote(filepath.Dir(remotePath)), shellQuote(temporary), shellQuote(temporary), shellQuote(artifact.SHA256), shellQuote(temporary), shellQuote(remotePath)) - command := session.remoteCommand(host, remote) + command := session.remoteCommand(host, remoteStageScript(remotePath, artifact.SHA256, "600", "checkpoint")) command.Stdin = file var output strings.Builder command.Stdout, command.Stderr = &output, &output diff --git a/internal/cli/model_hostfile_test.go b/internal/cli/model_hostfile_test.go index 02267821..71f3215f 100644 --- a/internal/cli/model_hostfile_test.go +++ b/internal/cli/model_hostfile_test.go @@ -346,8 +346,9 @@ case "$*" in exit 0 ;; *) - cat >/dev/null - exit 0 + while [ "$#" -gt 0 ] && [ "$1" != -- ]; do shift; done + shift 2 + exec sh -c "$1" ;; esac ` @@ -368,7 +369,7 @@ esac } t.Cleanup(func() { listenHostfileRendezvous = previousListener }) cluster := training.Cluster{Nodes: 2, Rendezvous: "127.0.0.1:0", RendezvousID: "session-test"} - scratch := t.TempDir() + scratch := privateTempDir(t) cacheRoot := t.TempDir() cache, err := lookaside.NewCache(cacheRoot, nil, lookaside.WithPersistentStorage(scratch, 1<<20)) if err != nil { @@ -388,6 +389,7 @@ esac if !strings.HasPrefix(session.remoteRoot, scratch+string(os.PathSeparator)) { t.Fatalf("remote launch root = %q, want beneath configured scratch %q", session.remoteRoot, scratch) } + assertFileSHA256(t, session.remoteBinary, session.binarySHA256) evaluation := training.EvaluationSet{Selection: "lowest-sha256-v1", SHA256: strings.Repeat("a", 64)} if err := session.publish(model.MultiNodePlan{ Kind: model.MultiNodePlanKind, Schema: model.MultiNodePlanSchema, @@ -752,3 +754,201 @@ func TestHostfileRemoteCommandUsesFuzzballWrapperContract(t *testing.T) { t.Fatal("Fuzzball wrapper lacks bounded graceful cancellation") } } + +// installLocalSSH installs an ssh shim that runs the remote script on the +// local host, exercising the real staging scripts with the test runner's +// tools. +func installLocalSSH(t *testing.T) { + t.Helper() + bin := t.TempDir() + script := "#!/bin/sh\nwhile [ \"$#\" -gt 0 ] && [ \"$1\" != -- ]; do shift; done\nshift 2\nexec sh -c \"$1\"\n" + if err := os.WriteFile(filepath.Join(bin, "ssh"), []byte(script), 0o755); err != nil { + t.Fatal(err) + } + t.Setenv("PATH", bin+string(os.PathListSeparator)+os.Getenv("PATH")) +} + +func newStagingSession(t *testing.T, scratch string) *hostfileSession { + t.Helper() + binary := filepath.Join(t.TempDir(), "waldo") + if err := os.WriteFile(binary, []byte("#!/bin/sh\nexit 0\n"), 0o700); err != nil { + t.Fatal(err) + } + digest, err := fileSHA256(binary) + if err != nil { + t.Fatal(err) + } + root := filepath.Join(scratch, "multinode", digest, "session-test") + return &hostfileSession{ + ctx: context.Background(), cancel: func() {}, + hostfile: trainingHostfile{Hosts: []string{"train-0", "train-1"}}, + binary: binary, binarySHA256: digest, + remoteRoot: root, remoteBinary: filepath.Join(root, "waldo"), resumeRoot: filepath.Join(root, "resume"), + output: io.Discard, + } +} + +// privateTempDir returns a test directory that staging accepts. The testing +// package creates t.TempDir with mode 0777 less the umask, which is +// group-writable under the common umask 002. +func privateTempDir(t *testing.T) string { + t.Helper() + directory := t.TempDir() + if err := os.Chmod(directory, 0o700); err != nil { + t.Fatal(err) + } + return directory +} + +func privateScratch(t *testing.T) string { + t.Helper() + scratch := filepath.Join(privateTempDir(t), "scratch") + if err := os.Mkdir(scratch, 0o700); err != nil { + t.Fatal(err) + } + return scratch +} + +func assertFileSHA256(t *testing.T, path, want string) { + t.Helper() + got, err := fileSHA256(path) + if err != nil { + t.Fatal(err) + } + if got != want { + t.Fatalf("%s SHA-256 = %s, want %s", path, got, want) + } +} + +func assertNoTransferFiles(t *testing.T, directory string) { + t.Helper() + entries, err := os.ReadDir(directory) + if err != nil && !os.IsNotExist(err) { + t.Fatal(err) + } + for _, entry := range entries { + if strings.HasPrefix(entry.Name(), ".waldo-transfer.") { + t.Fatalf("temporary upload %s remains in %s", entry.Name(), directory) + } + } +} + +func TestHostfileStageBinaryIntoVerifiedScratch(t *testing.T) { + installLocalSSH(t) + session := newStagingSession(t, privateScratch(t)) + for attempt := 0; attempt < 2; attempt++ { + if err := session.stageBinary("train-1"); err != nil { + t.Fatalf("stage attempt %d: %v", attempt+1, err) + } + } + assertFileSHA256(t, session.remoteBinary, session.binarySHA256) + for path, want := range map[string]os.FileMode{session.remoteBinary: 0o700, session.remoteRoot: 0o700} { + info, err := os.Stat(path) + if err != nil { + t.Fatal(err) + } + if info.Mode().Perm() != want { + t.Fatalf("%s mode = %v, want %v", path, info.Mode().Perm(), want) + } + } + assertNoTransferFiles(t, session.remoteRoot) + session.cleanupStaging() + if _, err := os.Stat(session.remoteRoot); !os.IsNotExist(err) { + t.Fatalf("launch staging remains after cleanup: %v", err) + } +} + +func TestHostfileStageBinaryRefusesDirectoriesOtherAccountsCanModify(t *testing.T) { + installLocalSSH(t) + for _, test := range []struct { + name string + mode os.FileMode + }{ + {name: "other-writable", mode: 0o777}, + {name: "group-writable", mode: 0o770}, + } { + t.Run(test.name, func(t *testing.T) { + scratch := privateScratch(t) + if err := os.Chmod(scratch, test.mode); err != nil { + t.Fatal(err) + } + session := newStagingSession(t, scratch) + err := session.stageBinary("train-1") + if err == nil || !strings.Contains(err.Error(), "writable by its group or other users") { + t.Fatalf("stage error = %v, want writable-directory refusal", err) + } + if !strings.Contains(err.Error(), filepath.Base(scratch)) { + t.Fatalf("stage error %q does not name the unsafe directory", err) + } + if _, statErr := os.Stat(session.remoteBinary); !os.IsNotExist(statErr) { + t.Fatalf("binary was staged in an unsafe directory: %v", statErr) + } + }) + } + t.Run("symlink-to-writable", func(t *testing.T) { + base := privateTempDir(t) + shared := filepath.Join(base, "shared") + if err := os.Mkdir(shared, 0o700); err != nil { + t.Fatal(err) + } + if err := os.Chmod(shared, 0o777); err != nil { + t.Fatal(err) + } + link := filepath.Join(base, "scratch") + if err := os.Symlink(shared, link); err != nil { + t.Fatal(err) + } + session := newStagingSession(t, link) + if err := session.stageBinary("train-1"); err == nil || !strings.Contains(err.Error(), "shared is writable") { + t.Fatalf("stage error = %v, want refusal of the resolved writable directory", err) + } + }) +} + +func TestHostfileStageBinaryFailsClosedOnDigestMismatch(t *testing.T) { + installLocalSSH(t) + session := newStagingSession(t, privateScratch(t)) + session.binarySHA256 = strings.Repeat("0", 64) + err := session.stageBinary("train-1") + if err == nil || !strings.Contains(err.Error(), "does not match") { + t.Fatalf("stage error = %v, want digest mismatch", err) + } + if _, statErr := os.Stat(session.remoteBinary); !os.IsNotExist(statErr) { + t.Fatalf("binary was published despite a digest mismatch: %v", statErr) + } + assertNoTransferFiles(t, session.remoteRoot) +} + +func TestHostfileStageResumeVerifiesCheckpointsInPrivateScratch(t *testing.T) { + installLocalSSH(t) + scratch := privateScratch(t) + session := newStagingSession(t, scratch) + source := filepath.Join(t.TempDir(), "runtime.pt") + if err := os.WriteFile(source, []byte("synthetic checkpoint"), 0o600); err != nil { + t.Fatal(err) + } + digest, err := fileSHA256(source) + if err != nil { + t.Fatal(err) + } + target := filepath.Join(session.resumeRoot, "run-42", "runtime.pt") + wrong := training.Artifact{Path: "runtime.pt", SHA256: strings.Repeat("0", 64)} + if err := session.stageResume("train-1", []training.Artifact{wrong}, []string{source}, []string{target}); err == nil || !strings.Contains(err.Error(), "does not match") { + t.Fatalf("stage error = %v, want digest mismatch", err) + } + if _, statErr := os.Stat(target); !os.IsNotExist(statErr) { + t.Fatalf("checkpoint was published despite a digest mismatch: %v", statErr) + } + assertNoTransferFiles(t, filepath.Dir(target)) + artifact := training.Artifact{Path: "runtime.pt", SHA256: digest} + if err := session.stageResume("train-1", []training.Artifact{artifact}, []string{source}, []string{target}); err != nil { + t.Fatal(err) + } + assertFileSHA256(t, target, digest) + if err := os.Chmod(scratch, 0o777); err != nil { + t.Fatal(err) + } + if err := session.stageResume("train-1", []training.Artifact{artifact}, []string{source}, []string{target}); err == nil || !strings.Contains(err.Error(), "writable by its group or other users") { + t.Fatalf("stage error = %v, want an already staged checkpoint in an unsafe directory to be refused", err) + } +}