Adding option to change the root directory name in build_tar.py and archive.py.

This is required to support Windows Docker image creation.

RELNOTES: None
PiperOrigin-RevId: 206326316

3 files changed, 87 insertions, 11 deletions

diff --git a/tools/build_defs/pkg/archive.py b/tools/build_defs/pkg/archive.py
index 16695232a4..c572c2d930 100644
--- a/tools/build_defs/pkg/archive.py
+++ b/tools/build_defs/pkg/archive.py
@@ -102,7 +102,7 @@ class TarFileWriter(object):
   class Error(Exception):
     pass
 
-  def __init__(self, name, compression=''):
+  def __init__(self, name, compression='', root_directory='./'):
     if compression in ['bzip2', 'bz2']:
       mode = 'w:bz2'
     else:
@@ -111,6 +111,8 @@ class TarFileWriter(object):
     # Support xz compression through xz... until we can use Py3
     self.xz = compression in ['xz', 'lzma']
     self.name = name
+    self.root_directory = root_directory.rstrip('/')
+
     self.fileobj = None
     if self.gz:
       # The Tarfile class doesn't allow us to specify gzip's mtime attribute.
@@ -155,8 +157,9 @@ class TarFileWriter(object):
       TarFileWriter.Error: when the recursion depth has exceeded the
                            `depth` argument.
     """
-    if not (name == '.' or name.startswith('/') or name.startswith('./')):
-      name = './' + name
+    if not (name == self.root_directory or name.startswith('/') or
+            name.startswith(self.root_directory + '/')):
+      name = os.path.join(self.root_directory, name)
     if os.path.isdir(path):
       # Remove trailing '/' (index -1 => last character)
       if name[-1] == '/':
@@ -238,8 +241,9 @@ class TarFileWriter(object):
       # Recurse into directory
       self.add_dir(name, file_content, uid, gid, uname, gname, mtime, mode)
       return
-    if not (name == '.' or name.startswith('/') or name.startswith('./')):
-      name = './' + name
+    if not (name == self.root_directory or name.startswith('/') or
+            name.startswith(self.root_directory + '/')):
+      name = os.path.join(self.root_directory, name)
     if kind == tarfile.DIRTYPE:
       name = name.rstrip('/')
       if name in self.directories:
@@ -300,7 +304,7 @@ class TarFileWriter(object):
       name_filter: filter out file by names. If not none, this method will be
           called for each file to add, given the name and should return true if
           the file is to be added to the final tar and false otherwise.
-      root: place all non-absolute content under given root direcory, if not
+      root: place all non-absolute content under given root directory, if not
           None.
 
     Raises:
@@ -351,8 +355,9 @@ class TarFileWriter(object):
           tarinfo.gname = ''
 
         name = tarinfo.name
-        if not name.startswith('/') and not name.startswith('.'):
-          name = './' + name
+        if (not name.startswith('/') and
+            not name.startswith(self.root_directory)):
+          name = os.path.join(self.root_directory, name)
         if root is not None:
           if name.startswith('.'):
             name = '.' + root + name.lstrip('.')
diff --git a/tools/build_defs/pkg/archive_test.py b/tools/build_defs/pkg/archive_test.py
index 5aa64079c3..2c52fbd423 100644
--- a/tools/build_defs/pkg/archive_test.py
+++ b/tools/build_defs/pkg/archive_test.py
@@ -295,5 +295,39 @@ class TarFileWriterTest(unittest.TestCase):
     ]
     self.assertTarFileContent(self.tempfile, content)
 
+  def testChangingRootDirectory(self):
+    with archive.TarFileWriter(self.tempfile, root_directory="root") as f:
+      f.add_file("d", tarfile.DIRTYPE)
+      f.add_file("d/f")
+
+      f.add_file("a", tarfile.DIRTYPE)
+      f.add_file("a/b", tarfile.DIRTYPE)
+      f.add_file("a/b", tarfile.DIRTYPE)
+      f.add_file("a/b/", tarfile.DIRTYPE)
+      f.add_file("a/b/c/f")
+
+      f.add_file("x/y/f")
+      f.add_file("x", tarfile.DIRTYPE)
+    content = [
+        {"name": "root",
+         "mode": 0o755},
+        {"name": "root/d",
+         "mode": 0o755},
+        {"name": "root/d/f"},
+        {"name": "root/a",
+         "mode": 0o755},
+        {"name": "root/a/b",
+         "mode": 0o755},
+        {"name": "root/a/b/c",
+         "mode": 0o755},
+        {"name": "root/a/b/c/f"},
+        {"name": "root/x",
+         "mode": 0o755},
+        {"name": "root/x/y",
+         "mode": 0o755},
+        {"name": "root/x/y/f"},
+    ]
+    self.assertTarFileContent(self.tempfile, content)
+
 if __name__ == "__main__":
   unittest.main()
diff --git a/tools/build_defs/pkg/build_tar.py b/tools/build_defs/pkg/build_tar.py
index e5428041a3..90d602d3d3 100644
--- a/tools/build_defs/pkg/build_tar.py
+++ b/tools/build_defs/pkg/build_tar.py
@@ -34,6 +34,11 @@ gflags.DEFINE_multistring('empty_file', [], 'An empty file to add to the layer')
 
 gflags.DEFINE_multistring('empty_dir', [], 'An empty dir to add to the layer')
 
+gflags.DEFINE_multistring(
+    'empty_root_dir',
+    [],
+    'An empty dir to add to the layer')
+
 gflags.DEFINE_multistring('tar', [], 'A tar file to add to the layer')
 
 gflags.DEFINE_multistring('deb', [], 'A debian package to add to the layer')
@@ -72,6 +77,9 @@ gflags.DEFINE_multistring('owner_names', None,
                           'Specify the owner names of individual files, e.g. '
                           'path/to/file=root.root.')
 
+gflags.DEFINE_string(
+    'root_directory', './', 'Default root directory is named "."')
+
 FLAGS = gflags.FLAGS
 
 
@@ -81,13 +89,18 @@ class TarFile(object):
   class DebError(Exception):
     pass
 
-  def __init__(self, output, directory, compression):
+  def __init__(self, output, directory, compression, root_directory):
     self.directory = directory
     self.output = output
     self.compression = compression
+    self.root_directory = root_directory
 
   def __enter__(self):
-    self.tarfile = archive.TarFileWriter(self.output, self.compression)
+    self.tarfile = archive.TarFileWriter(
+        self.output,
+        self.compression,
+        self.root_directory
+    )
     return self
 
   def __exit__(self, t, v, traceback):
@@ -175,6 +188,23 @@ class TarFile(object):
     self.add_empty_file(
         destpath, mode=mode, ids=ids, names=names, kind=tarfile.DIRTYPE)
 
+  def add_empty_root_dir(self, destpath, mode=None, ids=None, names=None):
+    """Add a directory to the root of the tar file.
+
+    Args:
+       destpath: the name of the directory in the layer
+       mode: force to set the specified mode, defaults to 644
+       ids: (uid, gid) for the file to set ownership
+       names: (username, groupname) for the file to set ownership.
+
+    An empty directory will be created as `destfile` in the root layer.
+    """
+    original_root_directory = self.tarfile.root_directory
+    self.tarfile.root_directory = destpath
+    self.add_empty_dir(
+        destpath, mode=mode, ids=ids, names=names)
+    self.tarfile.root_directory = original_root_directory
+
   def add_tar(self, tar):
     """Merge a tar file into the destination tar file.
 
@@ -265,7 +295,12 @@ def main(unused_argv):
       ids_map[f] = (int(user), int(group))
 
   # Add objects to the tar file
-  with TarFile(FLAGS.output, FLAGS.directory, FLAGS.compression) as output:
+  with TarFile(
+      FLAGS.output,
+      FLAGS.directory,
+      FLAGS.compression,
+      FLAGS.root_directory
+  ) as output:
 
     def file_attributes(filename):
       if filename[0] == '/':
@@ -283,6 +318,8 @@ def main(unused_argv):
       output.add_empty_file(f, **file_attributes(f))
     for f in FLAGS.empty_dir:
       output.add_empty_dir(f, **file_attributes(f))
+    for f in FLAGS.empty_root_dir:
+      output.add_empty_root_dir(f, **file_attributes(f))
     for tar in FLAGS.tar:
       output.add_tar(tar)
     for deb in FLAGS.deb: