fix: use atomic writes for state.json to prevent corruption

Signed-off-by: mayanksharmaCSE <mayanksharmacse1@gmail.com>

Author: MayankSharmaCSE

pkg/unikontainers/unikontainers.go

@@ -706,7 +706,7 @@ func (u *Unikontainer) saveContainerState() error {
 	}
 
 	stateName := filepath.Join(u.BaseDir, stateFilename)
-	return os.WriteFile(stateName, data, 0o644) //nolint: gosec
+	return atomicWriteFile(stateName, data, 0o644) //nolint: gosec
 }
 
 // getHooksByName returns the hooks for a given lifecycle stage

pkg/unikontainers/utils.go

@@ -106,24 +106,44 @@ func loadSpec(bundleDir string) (*specs.Spec, error) {
 	return &spec, nil
 }
 
-// writePidFile writes the content of pid to the file defined by path
-func writePidFile(path string, pid int) error {
-	var (
-		tmpDir  = filepath.Dir(path)
-		tmpName = filepath.Join(tmpDir, "."+filepath.Base(path))
-	)
-	f, err := os.OpenFile(tmpName, os.O_RDWR|os.O_CREATE|os.O_EXCL|os.O_SYNC, 0o666)
+// atomicWriteFile writes data to a file atomically by first writing to a
+// temporary file in the same directory, syncing it, and then renaming it
+// to the target path. This prevents partial/corrupt files if the process
+// is killed mid-write.
+func atomicWriteFile(path string, data []byte, perm os.FileMode) error {
+	dir := filepath.Dir(path)
+	tmpName := filepath.Join(dir, "."+filepath.Base(path)+".tmp")
+
+	f, err := os.OpenFile(tmpName, os.O_RDWR|os.O_CREATE|os.O_TRUNC|os.O_SYNC, perm)
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to create temp file: %w", err)
 	}
-	_, err = f.WriteString(strconv.Itoa(pid))
-	f.Close()
-	if err != nil {
-		return err
+
+	_, writeErr := f.Write(data)
+	syncErr := f.Sync()
+	closeErr := f.Close()
+
+	if writeErr != nil {
+		os.Remove(tmpName)
+		return fmt.Errorf("failed to write temp file: %w", writeErr)
 	}
+	if syncErr != nil {
+		os.Remove(tmpName)
+		return fmt.Errorf("failed to sync temp file: %w", syncErr)
+	}
+	if closeErr != nil {
+		os.Remove(tmpName)
+		return fmt.Errorf("failed to close temp file: %w", closeErr)
+	}
+
 	return os.Rename(tmpName, path)
 }
 
+// writePidFile writes the content of pid to the file defined by path
+func writePidFile(path string, pid int) error {
+	return atomicWriteFile(path, []byte(strconv.Itoa(pid)), 0o666)
+}
+
 // handleQueueProxy adds a hardcoded IP to the process's environment.
 // Then, the container is identified as a non-bima container
 // is spawned using runc.

pkg/unikontainers/utils_test.go

@@ -25,6 +25,60 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+func TestAtomicWriteFile(t *testing.T) {
+	t.Run("writes file atomically", func(t *testing.T) {
+		t.Parallel()
+		tmpDir := t.TempDir()
+		target := filepath.Join(tmpDir, "state.json")
+		data := []byte(`{"status":"running"}`)
+
+		err := atomicWriteFile(target, data, 0o644)
+		assert.NoError(t, err)
+
+		content, err := os.ReadFile(target)
+		assert.NoError(t, err)
+		assert.Equal(t, data, content)
+	})
+
+	t.Run("overwrites existing file", func(t *testing.T) {
+		t.Parallel()
+		tmpDir := t.TempDir()
+		target := filepath.Join(tmpDir, "state.json")
+
+		err := os.WriteFile(target, []byte("old"), 0o600)
+		assert.NoError(t, err)
+
+		newData := []byte("new content")
+		err = atomicWriteFile(target, newData, 0o644)
+		assert.NoError(t, err)
+
+		content, err := os.ReadFile(target)
+		assert.NoError(t, err)
+		assert.Equal(t, newData, content)
+	})
+
+	t.Run("no temp file left on success", func(t *testing.T) {
+		t.Parallel()
+		tmpDir := t.TempDir()
+		target := filepath.Join(tmpDir, "state.json")
+
+		err := atomicWriteFile(target, []byte("data"), 0o644)
+		assert.NoError(t, err)
+
+		tmpFile := filepath.Join(tmpDir, ".state.json.tmp")
+		_, err = os.Stat(tmpFile)
+		assert.True(t, os.IsNotExist(err), "Temp file should not exist after successful write")
+	})
+
+	t.Run("fails on invalid directory", func(t *testing.T) {
+		t.Parallel()
+		target := filepath.Join("/nonexistent/dir", "state.json")
+
+		err := atomicWriteFile(target, []byte("data"), 0o644)
+		assert.Error(t, err)
+	})
+}
+
 func TestWritePidFile(t *testing.T) {
 	tmpDir := t.TempDir() // Create a temporary directory for the test
 	pidFilePath := filepath.Join(tmpDir, "test.pid")