fix(agent/agentfiles): atomic writes and permission preservation (#23336)

Both writeFile and editFile now use the same atomic write strategy: temp file in the same directory, write, rename. This ensures a failed write leaves the original file intact instead of truncated. editFile already used temp-and-rename but lost the original file's permissions because afero.TempFile creates with mode 0600. Both functions now Chmod after rename to preserve the original mode. writeFile also swallowed io.Copy errors (logged but returned HTTP 200). Fixed to return the error so the client knows the write failed.
2026-09-22 05:05:20 +08:00 · 2026-03-20 01:56:19 +02:00
parent c2bc2c5738
commit de4e568994
2 changed files with 224 additions and 13 deletions
@@ -333,22 +333,68 @@ func (api *API) writeFile(ctx context.Context, r *http.Request, path string) (HT
 		return status, err
 	}

-	f, err := api.filesystem.Create(path)
+	// Check if the target already exists so we can preserve its
+	// permissions on the temp file before rename.
+	var origMode os.FileMode
+	var haveOrigMode bool
+	if stat, serr := api.filesystem.Stat(path); serr == nil {
+		if stat.IsDir() {
+			return http.StatusBadRequest, xerrors.Errorf("open %s: is a directory", path)
+		}
+		origMode = stat.Mode()
+		haveOrigMode = true
+	}
+
+	// Write to a temp file in the same directory so the rename is
+	// always on the same device (atomic).
+	tmpfile, err := afero.TempFile(api.filesystem, dir, filepath.Base(path))
 	if err != nil {
 		status := http.StatusInternalServerError
-		switch {
-		case errors.Is(err, os.ErrPermission):
+		if errors.Is(err, os.ErrPermission) {
 			status = http.StatusForbidden
-		case errors.Is(err, syscall.EISDIR):
-			status = http.StatusBadRequest
 		}
 		return status, err
 	}
-	defer f.Close()
+	tmpName := tmpfile.Name()

-	_, err = io.Copy(f, r.Body)
-	if err != nil && !errors.Is(err, io.EOF) && ctx.Err() == nil {
-		api.logger.Error(ctx, "workspace agent write file", slog.Error(err))
+	_, err = io.Copy(tmpfile, r.Body)
+	if err != nil && !errors.Is(err, io.EOF) {
+		_ = tmpfile.Close()
+		if rerr := api.filesystem.Remove(tmpName); rerr != nil {
+			api.logger.Warn(ctx, "unable to clean up temp file", slog.Error(rerr))
+		}
+		return http.StatusInternalServerError, xerrors.Errorf("write %s: %w", path, err)
+	}
+
+	// Close before rename to flush buffered data and catch write
+	// errors (e.g. delayed allocation failures).
+	if err := tmpfile.Close(); err != nil {
+		if rerr := api.filesystem.Remove(tmpName); rerr != nil {
+			api.logger.Warn(ctx, "unable to clean up temp file", slog.Error(rerr))
+		}
+		return http.StatusInternalServerError, xerrors.Errorf("write %s: %w", path, err)
+	}
+
+	// Set permissions on the temp file before rename so there is
+	// no window where the target has wrong permissions.
+	if haveOrigMode {
+		if err := api.filesystem.Chmod(tmpName, origMode); err != nil {
+			api.logger.Warn(ctx, "unable to set file permissions",
+				slog.F("path", path),
+				slog.Error(err),
+			)
+		}
+	}
+
+	if err := api.filesystem.Rename(tmpName, path); err != nil {
+		if rerr := api.filesystem.Remove(tmpName); rerr != nil {
+			api.logger.Warn(ctx, "unable to clean up temp file", slog.Error(rerr))
+		}
+		status := http.StatusInternalServerError
+		if errors.Is(err, os.ErrPermission) {
+			status = http.StatusForbidden
+		}
+		return status, err
 	}

 	return 0, nil
@@ -460,18 +506,44 @@ func (api *API) editFile(ctx context.Context, path string, edits []workspacesdk.
 	if err != nil {
 		return http.StatusInternalServerError, err
 	}
-	defer tmpfile.Close()
+	tmpName := tmpfile.Name()

 	if _, err := tmpfile.Write([]byte(content)); err != nil {
-		if rerr := api.filesystem.Remove(tmpfile.Name()); rerr != nil {
+		_ = tmpfile.Close()
+		if rerr := api.filesystem.Remove(tmpName); rerr != nil {
 			api.logger.Warn(ctx, "unable to clean up temp file", slog.Error(rerr))
 		}
 		return http.StatusInternalServerError, xerrors.Errorf("edit %s: %w", path, err)
 	}

-	err = api.filesystem.Rename(tmpfile.Name(), path)
+	// Close before rename to flush buffered data and catch write
+	// errors (e.g. delayed allocation failures).
+	if err := tmpfile.Close(); err != nil {
+		if rerr := api.filesystem.Remove(tmpName); rerr != nil {
+			api.logger.Warn(ctx, "unable to clean up temp file", slog.Error(rerr))
+		}
+		return http.StatusInternalServerError, xerrors.Errorf("edit %s: %w", path, err)
+	}
+
+	// Set permissions on the temp file before rename so there is
+	// no window where the target has wrong permissions.
+	if err := api.filesystem.Chmod(tmpName, stat.Mode()); err != nil {
+		api.logger.Warn(ctx, "unable to set file permissions",
+			slog.F("path", path),
+			slog.Error(err),
+		)
+	}
+
+	err = api.filesystem.Rename(tmpName, path)
 	if err != nil {
-		return http.StatusInternalServerError, err
+		if rerr := api.filesystem.Remove(tmpName); rerr != nil {
+			api.logger.Warn(ctx, "unable to clean up temp file", slog.Error(rerr))
+		}
+		status := http.StatusInternalServerError
+		if errors.Is(err, os.ErrPermission) {
+			status = http.StatusForbidden
+		}
+		return status, err
 	}

 	return 0, nil
@@ -14,6 +14,7 @@ import (
 	"strings"
 	"syscall"
 	"testing"
+	"testing/iotest"

 	"github.com/go-chi/chi/v5"
 	"github.com/google/uuid"
@@ -399,6 +400,83 @@ func TestWriteFile(t *testing.T) {
 	}
 }

+func TestWriteFile_ReportsIOError(t *testing.T) {
+	t.Parallel()
+
+	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
+	fs := afero.NewMemMapFs()
+	api := agentfiles.NewAPI(logger, fs, nil)
+
+	tmpdir := os.TempDir()
+	path := filepath.Join(tmpdir, "write-io-error")
+	err := afero.WriteFile(fs, path, []byte("original"), 0o644)
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel()
+
+	// A reader that always errors simulates a failed body read
+	// (e.g. network interruption). The atomic write should leave
+	// the original file intact.
+	body := iotest.ErrReader(xerrors.New("simulated I/O error"))
+	w := httptest.NewRecorder()
+	r := httptest.NewRequestWithContext(ctx, http.MethodPost,
+		fmt.Sprintf("/write-file?path=%s", path), body)
+	api.Routes().ServeHTTP(w, r)
+
+	require.Equal(t, http.StatusInternalServerError, w.Code)
+	got := &codersdk.Error{}
+	err = json.NewDecoder(w.Body).Decode(got)
+	require.NoError(t, err)
+	require.ErrorContains(t, got, "simulated I/O error")
+
+	// The original file must survive the failed write.
+	data, err := afero.ReadFile(fs, path)
+	require.NoError(t, err)
+	require.Equal(t, "original", string(data))
+}
+
+func TestWriteFile_PreservesPermissions(t *testing.T) {
+	t.Parallel()
+
+	if runtime.GOOS == "windows" {
+		t.Skip("file permissions are not reliably supported on Windows")
+	}
+
+	dir := t.TempDir()
+	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+	osFs := afero.NewOsFs()
+	api := agentfiles.NewAPI(logger, osFs, nil)
+
+	path := filepath.Join(dir, "script.sh")
+	err := afero.WriteFile(osFs, path, []byte("#!/bin/sh\necho hello\n"), 0o755)
+	require.NoError(t, err)
+
+	info, err := osFs.Stat(path)
+	require.NoError(t, err)
+	require.Equal(t, os.FileMode(0o755), info.Mode().Perm())
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel()
+
+	// Overwrite the file with new content.
+	w := httptest.NewRecorder()
+	r := httptest.NewRequestWithContext(ctx, http.MethodPost,
+		fmt.Sprintf("/write-file?path=%s", path),
+		bytes.NewReader([]byte("#!/bin/sh\necho world\n")))
+	api.Routes().ServeHTTP(w, r)
+	require.Equal(t, http.StatusOK, w.Code)
+
+	data, err := afero.ReadFile(osFs, path)
+	require.NoError(t, err)
+	require.Equal(t, "#!/bin/sh\necho world\n", string(data))
+
+	info, err = osFs.Stat(path)
+	require.NoError(t, err)
+	require.Equal(t, os.FileMode(0o755), info.Mode().Perm(),
+		"write_file should preserve the original file's permissions")
+}
+
 func TestEditFiles(t *testing.T) {
 	t.Parallel()

@@ -907,6 +985,67 @@ func TestEditFiles(t *testing.T) {
 	}
 }

+func TestEditFiles_PreservesPermissions(t *testing.T) {
+	t.Parallel()
+
+	if runtime.GOOS == "windows" {
+		t.Skip("file permissions are not reliably supported on Windows")
+	}
+
+	dir := t.TempDir()
+	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+	osFs := afero.NewOsFs()
+	api := agentfiles.NewAPI(logger, osFs, nil)
+
+	path := filepath.Join(dir, "script.sh")
+	err := afero.WriteFile(osFs, path, []byte("#!/bin/sh\necho hello\n"), 0o755)
+	require.NoError(t, err)
+
+	// Sanity-check the initial mode.
+	info, err := osFs.Stat(path)
+	require.NoError(t, err)
+	require.Equal(t, os.FileMode(0o755), info.Mode().Perm())
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel()
+
+	body := workspacesdk.FileEditRequest{
+		Files: []workspacesdk.FileEdits{
+			{
+				Path: path,
+				Edits: []workspacesdk.FileEdit{
+					{
+						Search:  "hello",
+						Replace: "world",
+					},
+				},
+			},
+		},
+	}
+	buf := bytes.NewBuffer(nil)
+	enc := json.NewEncoder(buf)
+	enc.SetEscapeHTML(false)
+	err = enc.Encode(body)
+	require.NoError(t, err)
+
+	w := httptest.NewRecorder()
+	r := httptest.NewRequestWithContext(ctx, http.MethodPost, "/edit-files", buf)
+	api.Routes().ServeHTTP(w, r)
+	require.Equal(t, http.StatusOK, w.Code)
+
+	// Verify content was updated.
+	data, err := afero.ReadFile(osFs, path)
+	require.NoError(t, err)
+	require.Equal(t, "#!/bin/sh\necho world\n", string(data))
+
+	// Verify permissions are preserved after the
+	// temp-file-and-rename cycle.
+	info, err = osFs.Stat(path)
+	require.NoError(t, err)
+	require.Equal(t, os.FileMode(0o755), info.Mode().Perm(),
+		"edit_files should preserve the original file's permissions")
+}
+
 func TestHandleWriteFile_ChatHeaders_UpdatesPathStore(t *testing.T) {
 	t.Parallel()