Repository: canonical/operator
Pull request #574 (2021-07-23 05:46:43)

Fix Pebble push() handling of binary files
Python's mime code's msg.as_bytes() ends up [here](https://github.com/python/cpython/blob/0f55d21212a0fb6dec4479eb76249ab2b54e57a3/Lib/email/generator.py#L149-L163
). But that translates `\r\n` into `\n`, which is obviously horrible for binary files. I saw someone on StackOverflow that had subclassed Generator and overriden this private _write_lines method to avoid this, but that seems like a total hack. Besides, it's not hard to just do this ourselves and get control over it, so let's do that.

I've added regression tests for this (and verified the `push` tests fail before the fix), and I've also tested manually using `test/pebble_cli.py` against the real Pebble with both a small file with mixed line endings in it and a huge binary executable (again, fails before the fix, works fine after). The `pull` function doesn't have this problem.

Fixes https://github.com/canonical/operator/issues/573.

jnsgruk: Nice work @benhoyt!

Patch:
diff --git a/ops/pebble.py b/ops/pebble.py
index 5e13fd3..c65e8c9 100644
--- a/ops/pebble.py
+++ b/ops/pebble.py
@@ -17,7 +17,7 @@
 For a command-line interface for local testing, see test/pebble_cli.py.
 """
 
-from email.mime.multipart import MIMEBase, MIMEMultipart
+import binascii
 import cgi
 import datetime
 import email.parser
@@ -25,6 +25,7 @@ import enum
 import http.client
 import io
 import json
+import os
 import re
 import socket
 import sys
@@ -1011,28 +1012,18 @@ class Client:
             'files': [info],
         }
 
-        multipart = MIMEMultipart('form-data')
-
-        part = MIMEBase('application', 'json')
-        part.add_header('Content-Disposition', 'form-data', name='request')
-        part.set_payload(json.dumps(metadata))
-        multipart.attach(part)
-
-        part = MIMEBase('application', 'octet-stream')
-        part.add_header('Content-Disposition', 'form-data', name='files', filename=path)
         if hasattr(source, 'read'):
             content = source.read()
         else:
             content = source
         if isinstance(content, str):
             content = content.encode(encoding)
-        part.set_payload(content)
-        multipart.attach(part)
 
-        data = multipart.as_bytes()  # must be called before accessing multipart['Content-Type']
+        data, content_type = self._encode_multipart(metadata, path, content)
+
         headers = {
             'Accept': 'application/json',
-            'Content-Type': multipart['Content-Type'],
+            'Content-Type': content_type,
         }
         response = self._request_raw('POST', '/v1/files', None, headers, data)
         self._ensure_content_type(response.headers, 'application/json')
@@ -1054,6 +1045,29 @@ class Client:
             d['group'] = group
         return d
 
+    @staticmethod
+    def _encode_multipart(metadata, path, content):
+        # Python's stdlib mime/multipart handling is screwy and doesn't handle
+        # binary properly, so roll our own.
+        boundary = binascii.hexlify(os.urandom(16))
+        path_escaped = path.replace('"', '\\"').encode('utf-8')  # NOQA: test_quote_backslashes
+        parts = []
+        parts.extend([
+            b'--', boundary, b'\r\n',
+            b'Content-Type: application/json\r\n',
+            b'Content-Disposition: form-data; name="request"\r\n',
+            b'\r\n',
+            json.dumps(metadata).encode('utf-8'), b'\r\n',
+            b'--', boundary, b'\r\n',
+            b'Content-Type: application/octet-stream\r\n',
+            b'Content-Disposition: form-data; name="files"; filename="', path_escaped, b'"\r\n',
+            b'\r\n',
+            content, b'\r\n',
+            b'--', boundary, b'--\r\n',
+        ])
+        content_type = 'multipart/form-data; boundary="' + boundary.decode('utf-8') + '"'
+        return b''.join(parts), content_type
+
     def list_files(self, path: str, *, pattern: str = None,
                    itself: bool = False) -> typing.List[FileInfo]:
         """Return list of file information from given path on remote system.


Test patch:
diff --git a/test/test_pebble.py b/test/test_pebble.py
index 78f3c2d..9123864 100644
--- a/test/test_pebble.py
+++ b/test/test_pebble.py
@@ -1206,7 +1206,7 @@ services:
 --01234567890123456789012345678901
 Content-Disposition: form-data; name="files"; filename="/etc/hosts"
 
-127.0.0.1 localhost  # """ + b'\xf0\x9f\x98\x80' + b"""
+127.0.0.1 localhost  # """ + b'\xf0\x9f\x98\x80\nfoo\r\nbar' + b"""
 --01234567890123456789012345678901
 Content-Disposition: form-data; name="response"
 
@@ -1221,7 +1221,7 @@ Content-Disposition: form-data; name="response"
         ))
 
         content = self.client.pull('/etc/hosts').read()
-        self.assertEqual(content, '127.0.0.1 localhost  # 😀')
+        self.assertEqual(content, '127.0.0.1 localhost  # 😀\nfoo\r\nbar')
 
         self.assertEqual(self.client.requests, [
             ('GET', '/v1/files', {'action': 'read', 'path': '/etc/hosts'},
@@ -1235,7 +1235,7 @@ Content-Disposition: form-data; name="response"
 --01234567890123456789012345678901
 Content-Disposition: form-data; name="files"; filename="/etc/hosts"
 
-127.0.0.1 localhost  # """ + b'\xf0\x9f\x98\x80' + b"""
+127.0.0.1 localhost  # """ + b'\xf0\x9f\x98\x80\nfoo\r\nbar' + b"""
 --01234567890123456789012345678901
 Content-Disposition: form-data; name="response"
 
@@ -1250,7 +1250,7 @@ Content-Disposition: form-data; name="response"
         ))
 
         content = self.client.pull('/etc/hosts', encoding=None).read()
-        self.assertEqual(content, b'127.0.0.1 localhost  # \xf0\x9f\x98\x80')
+        self.assertEqual(content, b'127.0.0.1 localhost  # \xf0\x9f\x98\x80\nfoo\r\nbar')
 
         self.assertEqual(self.client.requests, [
             ('GET', '/v1/files', {'action': 'read', 'path': '/etc/hosts'},
@@ -1329,10 +1329,10 @@ bad path
         self.assertEqual(str(cm.exception), 'no "response" field in multipart body')
 
     def test_push_str(self):
-        self._test_push_str('content 😀')
+        self._test_push_str('content 😀\nfoo\r\nbar')
 
     def test_push_text(self):
-        self._test_push_str(io.StringIO('content 😀'))
+        self._test_push_str(io.StringIO('content 😀\nfoo\r\nbar'))
 
     def _test_push_str(self, source):
         self.client.responses.append((
@@ -1359,17 +1359,17 @@ bad path
         content_type = headers['Content-Type']
         req, filename, content = self._parse_write_multipart(content_type, body)
         self.assertEqual(filename, '/foo/bar')
-        self.assertEqual(content, b'content \xf0\x9f\x98\x80')
+        self.assertEqual(content, b'content \xf0\x9f\x98\x80\nfoo\r\nbar')
         self.assertEqual(req, {
             'action': 'write',
             'files': [{'path': '/foo/bar'}],
         })
 
     def test_push_bytes(self):
-        self._test_push_bytes(b'content \xf0\x9f\x98\x80')
+        self._test_push_bytes(b'content \xf0\x9f\x98\x80\nfoo\r\nbar')
 
     def test_push_binary(self):
-        self._test_push_bytes(io.BytesIO(b'content \xf0\x9f\x98\x80'))
+        self._test_push_bytes(io.BytesIO(b'content \xf0\x9f\x98\x80\nfoo\r\nbar'))
 
     def _test_push_bytes(self, source):
         self.client.responses.append((
@@ -1396,7 +1396,7 @@ bad path
         content_type = headers['Content-Type']
         req, filename, content = self._parse_write_multipart(content_type, body)
         self.assertEqual(filename, '/foo/bar')
-        self.assertEqual(content, b'content \xf0\x9f\x98\x80')
+        self.assertEqual(content, b'content \xf0\x9f\x98\x80\nfoo\r\nbar')
         self.assertEqual(req, {
             'action': 'write',
             'files': [{'path': '/foo/bar'}],