Make pygrub cope better with big files in the guest.
Only read the first megabyte of a configuration file (grub etc.) and read
the kernel and ramdisk files from the guest in one megabyte pieces
so pygrub doesn't use a lot of memory if the files are large.
With --not-really option check that the chosen kernel and ramdisk files exist.
If there are problems writing the copy of the kernel or ramdisk, delete the
copied files and exit in case they have filled the filesystem.

Signed-off-by: Michael Young <m.a.young@xxxxxxxxxxxx>

--- xen-4.2.0/tools/pygrub/src/pygrub.orig	2012-05-12 16:40:48.000000000 +0100
+++ xen-4.2.0/tools/pygrub/src/pygrub	2012-06-25 21:53:49.556446369 +0100
@@ -28,6 +28,7 @@
 import grub.ExtLinuxConf
 
 PYGRUB_VER = 0.6
+fs_read_max=1048576
 
 def enable_cursor(ison):
     if ison:
@@ -448,7 +449,8 @@
         if self.__dict__.get('cf', None) is None:
             raise RuntimeError, "couldn't find bootloader config file in the image provided."
         f = fs.open_file(self.cf.filename)
-        buf = f.read()
+        # limit read size to avoid pathological cases
+        buf = f.read(fs_read_max)
         del f
         self.cf.parse(buf)
 
@@ -697,6 +699,39 @@
     def usage():
         print >> sys.stderr, "Usage: %s [-q|--quiet] [-i|--interactive] [-n|--not-really] [--output=] [--kernel=] [--ramdisk=] [--args=] [--entry=] [--output-directory=] [--output-format=sxp|simple|simple0] <image>" %(sys.argv[0],)
 
+    def copy_from_image(fs, file_to_read, file_type, output_directory,
+                                              not_really, clean_extra_file):
+        if not_really:
+            if fs.file_exists(file_to_read):
+                return "<%s:%s>" % (file_type,file_to_read)
+            else:
+                sys.exit("The requested %s file does not exist" % file_type)
+        try:
+            datafile = fs.open_file(file_to_read)
+        except:
+            if clean_extra_file:
+                os.unlink(clean_extra_file)
+            sys.exit("Error opening %s" % file_to_read)
+        (tfd, ret) = tempfile.mkstemp(prefix="boot_"+file_type+".",
+                                                       dir=output_directory)
+        dataoff=0
+        while True:
+            data=datafile.read(fs_read_max,dataoff)
+            if len(data) == 0:
+                os.close(tfd)
+                del datafile
+                return ret
+            try:
+                os.write(tfd, data)
+            except:
+                os.close(tfd)
+                os.unlink(ret)
+                del datafile
+                if clean_extra_file:
+                    os.unlink(clean_extra_file)
+                sys.exit("error writing temporary copy of "+file_type)
+            dataoff+=len(data)
+
     try:
         opts, args = getopt.gnu_getopt(sys.argv[1:], 'qinh::',
                                    ["quiet", "interactive", "not-really", "help", 
@@ -821,24 +856,12 @@
     if not fs:
         raise RuntimeError, "Unable to find partition containing kernel"
 
-    if not_really:
-        bootcfg["kernel"] = "<kernel:%s>" % chosencfg["kernel"]
-    else:
-        data = fs.open_file(chosencfg["kernel"]).read()
-        (tfd, bootcfg["kernel"]) = tempfile.mkstemp(prefix="boot_kernel.",
-                                                    dir=output_directory)
-        os.write(tfd, data)
-        os.close(tfd)
+    bootcfg["kernel"] = copy_from_image(fs, chosencfg["kernel"], 
+                              "kernel", output_directory, not_really, "")
 
     if chosencfg["ramdisk"]:
-        if not_really:
-            bootcfg["ramdisk"] = "<ramdisk:%s>" % chosencfg["ramdisk"]
-        else:
-            data = fs.open_file(chosencfg["ramdisk"],).read()
-            (tfd, bootcfg["ramdisk"]) = tempfile.mkstemp(
-                prefix="boot_ramdisk.", dir=output_directory)
-            os.write(tfd, data)
-            os.close(tfd)
+        bootcfg["ramdisk"] = copy_from_image(fs, chosencfg["ramdisk"],
+              "ramdisk", output_directory, not_really, bootcfg["kernel"])
     else:
         initrd = None