Create FileRecoveryService to restore corrupted files after a crash

atom · as-cii · May 27, 2016 · May 23, 2016 · May 23, 2016 · May 23, 2016
commit b58ce49d0d38381265de9e0af3a9315ccb95f3ba
diff --git a/src/browser/atom-application.coffee b/src/browser/atom-application.coffee
@@ -4,6 +4,7 @@ AtomProtocolHandler = require './atom-protocol-handler'
 AutoUpdateManager = require './auto-update-manager'
 StorageFolder = require '../storage-folder'
 Config = require '../config'
+FileRecoveryService = require './file-recovery-service'
 ipcHelpers = require '../ipc-helpers'
 {BrowserWindow, Menu, app, dialog, ipcMain, shell} = require 'electron'
 fs = require 'fs-plus'
@@ -78,12 +79,14 @@ class AtomApplication
     @autoUpdateManager = new AutoUpdateManager(@version, options.test, @resourcePath, @config)
     @applicationMenu = new ApplicationMenu(@version, @autoUpdateManager)
     @atomProtocolHandler = new AtomProtocolHandler(@resourcePath, @safeMode)
+    @fileRecoveryService = new FileRecoveryService(path.join(process.env.ATOM_HOME, "recovery"))
 
     @listenForArgumentsFromNewProcess()
     @setupJavaScriptArguments()
     @handleEvents()
     @setupDockMenu()
     @storageFolder = new StorageFolder(process.env.ATOM_HOME)
+    @fileRecoveryService.start()
 
     if options.pathsToOpen?.length > 0 or options.urlsToOpen?.length > 0 or options.test
       @openWithOptions(options)

diff --git a/src/browser/file-recovery-service.js b/src/browser/file-recovery-service.js
@@ -0,0 +1,60 @@
+'use babel'
+
+import {ipcMain} from 'electron'
+import crypto from 'crypto'
+import Path from 'path'
+import fs from 'fs-plus'
+
+export default class FileRecoveryService {
+  constructor (recoveryDirectory) {
+    this.recoveryDirectory = recoveryDirectory
+    this.recoveryPathsByWindowAndFilePath = new WeakMap
+    this.crashListeners = new WeakSet
+  }
+
+  start () {
+    this.willSavePathListener = ipcMain.on('will-save-path', this.willSavePath.bind(this))
+    this.didSavePathListener = ipcMain.on('did-save-path', this.didSavePath.bind(this))
+  }
+
+  willSavePath (event, path) {
+    if (!fs.existsSync(path)) {
+      // Unexisting files won't be truncated/overwritten, and so there's no data to be lost.
+      return
+    }
+
+    const window = event.sender
+    const recoveryFileName = crypto.createHash('sha1').update(path + Date.now().toString(), 'utf8').digest('hex').substring(0, 10)
+    const recoveryPath = Path.join(this.recoveryDirectory, recoveryFileName)
+    fs.writeFileSync(recoveryPath, fs.readFileSync(path))
+
+    if (!this.recoveryPathsByWindowAndFilePath.has(window)) {
+      this.recoveryPathsByWindowAndFilePath.set(window, new Map)
+    }
+    this.recoveryPathsByWindowAndFilePath.get(window).set(path, recoveryPath)
+
+    if (!this.crashListeners.has(window)) {
+      window.on('crashed', () => this.recoverFilesForWindow(window))
+      this.crashListeners.add(window)
+    }
+  }
+
+  didSavePath (event, path) {
+    const recoveryPathsByFilePath = this.recoveryPathsByWindowAndFilePath.get(event.sender)
+    if (recoveryPathsByFilePath.has(path)) {
+      const recoveryPath = recoveryPathsByFilePath.get(path)
+      fs.unlinkSync(recoveryPath)
+      recoveryPathsByFilePath.delete(path)
+    }
+  }
+
+  recoverFilesForWindow (window) {
+    const recoveryPathsByFilePath = this.recoveryPathsByWindowAndFilePath.get(window)
+    for (let [filePath, recoveryPath] of recoveryPathsByFilePath) {
+      fs.writeFileSync(filePath, fs.readFileSync(recoveryPath))
+      fs.unlinkSync(recoveryPath)
+    }
+
+    recoveryPathsByFilePath.clear()
+  }
+}