Make it harder to leak temp files on devices

build/android/pylib/android_commands.py

 # Copyright (c) 2012 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Provides an interface to communicate with the device via the adb command.
 
 Assumes adb binary is currently on system path.
 """
 # pylint: disable-all
 
 import collections
 import datetime
 import inspect
 import logging
 import os
+import random
 import re
 import shlex
 import signal
 import subprocess
 import sys
 import tempfile
 import time
 
 import cmd_helper
 import constants
@@ skipping 45 matching lines @@
 CONTROL_USB_CHARGING_COMMANDS = [
   {
     # Nexus 4
     'witness_file': '/sys/module/pm8921_charger/parameters/disabled',
     'enable_command': 'echo 0 > /sys/module/pm8921_charger/parameters/disabled',
     'disable_command':
         'echo 1 > /sys/module/pm8921_charger/parameters/disabled',
   },
 ]
 
+class DeviceTempFile(object):
+  def __init__(self, android_commands, prefix='temp_file', suffix=''):
+    """Find an unused temporary file path in the devices external directory.
+
+    When this object is closed, the file will be deleted on the device.
+    """
+    self.android_commands = android_commands
+    while True:
+      # TODO(cjhopman): This could actually return the same file in multiple
+      # calls if the caller doesn't write to the files immediately. This is
+      # expected to never happen.
+      i = random.randint(0, 1000000)
+      self.name = '%s/%s-%d-%010d%s' % (
+          android_commands.GetExternalStorage(),
+          prefix, int(time.time()), i, suffix)
+      if not android_commands.FileExistsOnDevice(self.name):
+        break
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, type, value, traceback):
+    self.close()
+
+  def close(self):
+    self.android_commands.RunShellCommand('rm ' + self.name)
+
+
 def GetAVDs():
   """Returns a list of AVDs."""
   re_avd = re.compile('^[ ]+Name: ([a-zA-Z0-9_:.-]+)', re.MULTILINE)
   avds = re_avd.findall(cmd_helper.GetCmdOutput(['android', 'list', 'avd']))
   return avds
 
 def ResetBadDevices():
   """Removes the blacklist that keeps track of bad devices for a current
      build.
   """
@@ skipping 1027 matching lines @@
     return self.RunShellCommand('cat "%s" 2>/dev/null' % filename,
                                 log_result=log_result)
 
   def SetFileContents(self, filename, contents):
     """Writes |contents| to the file specified by |filename|."""
     with tempfile.NamedTemporaryFile() as f:
       f.write(contents)
       f.flush()
       self._adb.Push(f.name, filename)
 
-  _TEMP_FILE_BASE_FMT = 'temp_file_%d'
-  _TEMP_SCRIPT_FILE_BASE_FMT = 'temp_script_file_%d.sh'
-
-  def _GetDeviceTempFileName(self, base_name):
-    i = 0
-    while self.FileExistsOnDevice(
-        self.GetExternalStorage() + '/' + base_name % i):
-      i += 1
-    return self.GetExternalStorage() + '/' + base_name % i
-
   def RunShellCommandWithSU(self, command, timeout_time=20, log_result=False):
     return self.RunShellCommand('su -c %s' % command, timeout_time, log_result)
 
   def CanAccessProtectedFileContents(self):
     """Returns True if Get/SetProtectedFileContents would work via "su" or adb
     shell running as root.
 
     Devices running user builds don't have adb root, but may provide "su" which
     can be used for accessing protected files.
     """
@@ skipping 37 matching lines @@
       return command_runner(command)
     else:
       logging.warning('Could not access protected file: %s' % filename)
       return []
 
   def SetProtectedFileContents(self, filename, contents):
     """Writes |contents| to the protected file specified by |filename|.
 
     This is less efficient than SetFileContents.
     """
-    temp_file = self._GetDeviceTempFileName(AndroidCommands._TEMP_FILE_BASE_FMT)
-    temp_script = self._GetDeviceTempFileName(
-        AndroidCommands._TEMP_SCRIPT_FILE_BASE_FMT)
+    with DeviceTempFile(self) as temp_file:
+      with DeviceTempFile(self, suffix=".sh") as temp_script:
+        # Put the contents in a temporary file
+        self.SetFileContents(temp_file.name, contents)
+        # Create a script to copy the file contents to its final destination
+        self.SetFileContents(temp_script.name,
+                             'cat %s > %s' % (temp_file.name, filename))
 
-    try:
-      # Put the contents in a temporary file
-      self.SetFileContents(temp_file, contents)
-      # Create a script to copy the file contents to its final destination
-      self.SetFileContents(temp_script, 'cat %s > %s' % (temp_file, filename))
+        command = 'sh %s' % temp_script.name
+        command_runner = self._GetProtectedFileCommandRunner()
+        if command_runner:
+          return command_runner(command)
+        else:
+          logging.warning(
+              'Could not set contents of protected file: %s' % filename)
 
-      command = 'sh %s' % temp_script
-      command_runner = self._GetProtectedFileCommandRunner()
-      if command_runner:
-        return command_runner(command)
-      else:
-        logging.warning(
-            'Could not set contents of protected file: %s' % filename)
-    finally:
-      # And remove the temporary files
-      self.RunShellCommand('rm ' + temp_file)
-      self.RunShellCommand('rm ' + temp_script)
 
   def RemovePushedFiles(self):
     """Removes all files pushed with PushIfNeeded() from the device."""
     for p in self._pushed_files:
       self.RunShellCommand('rm -r %s' % p, timeout_time=2 * 60)
 
   def ListPathContents(self, path):
     """Lists files in all subdirectories of |path|.
 
     Args:
@@ skipping 644 matching lines @@
 
     Uses a shell script running on the target to copy new and changed files the
     source directory to the destination directory and remove added files. This
     is in some cases much faster than cp -r.
 
     Args:
       source: absolute path of source directory
       dest: absolute path of destination directory
     """
     logging.info('In EfficientDeviceDirectoryCopy %s %s', source, dest)
-    temp_script_file = self._GetDeviceTempFileName(
-        AndroidCommands._TEMP_SCRIPT_FILE_BASE_FMT)
-    host_script_path = os.path.join(constants.DIR_SOURCE_ROOT,
-                                    'build',
-                                    'android',
-                                    'pylib',
-                                    'efficient_android_directory_copy.sh')
-    try:
-      self._adb.Push(host_script_path, temp_script_file)
-      self.EnableAdbRoot()
-      out = self.RunShellCommand('sh %s %s %s' % (temp_script_file,
-                                                  source,
-                                                  dest),
-                                 timeout_time=120)
+    with DeviceTempFile(self, suffix=".sh") as temp_script_file:
+      host_script_path = os.path.join(constants.DIR_SOURCE_ROOT,
+                                      'build',
+                                      'android',
+                                      'pylib',
+                                      'efficient_android_directory_copy.sh')
+      self._adb.Push(host_script_path, temp_script_file.name)
+      self.EnableAdbRoot
+      out = self.RunShellCommand(
+          'sh %s %s %s' % (temp_script_file.name, source, dest),
+          timeout_time=120)
       if self._device:
         device_repr = self._device[-4:]
       else:
         device_repr = '????'
       for line in out:
         logging.info('[%s]> %s', device_repr, line)
-    finally:
-      self.RunShellCommand('rm %s' % temp_script_file)
 
   def _GetControlUsbChargingCommand(self):
     if self._control_usb_charging_command['cached']:
       return self._control_usb_charging_command['command']
     self._control_usb_charging_command['cached'] = True
     for command in CONTROL_USB_CHARGING_COMMANDS:
       # Assert command is valid.
       assert 'disable_command' in command
       assert 'enable_command' in command
       assert 'witness_file' in command
@@ skipping 48 matching lines @@
   """
   def __init__(self, output):
     self._output = output
 
   def write(self, data):
     data = data.replace('\r\r\n', '\n')
     self._output.write(data)
 
   def flush(self):
     self._output.flush()