refactor: download and decompress xz images (#29)

* refactor: download and decompress xz images * fix: (intermediate value).getReader is not a function * test: use image worker to unpack and validate images * fix: use manifest defined size to calculate unpacking progress * fix: use system alt image size * fix: await writable.write promises so writes are complete before closing * fix: switch from xzwasm to xz-decompress for unpacking xz images xz-decompress is a fork of xzwasm that contains a fix for an issue causing checksum mismatches when awaiting writable writes --------- Co-authored-by: Justin Newberry <[email protected]>
commaai · Jan 28, 2024 · 398d72a · 398d72a
1 parent 3232596
commit 398d72a
Show file tree

Hide file tree

Showing 5 changed files with 76 additions and 67 deletions.
diff --git a/package.json b/package.json
@@ -21,12 +21,12 @@
     "jssha": "^3.3.1",
     "next": "13.4.1",
     "next-plausible": "^3.10.1",
-    "pako": "2.1.0",
     "postcss": "8.4.24",
     "react": "18.2.0",
     "react-dom": "18.2.0",
     "tailwindcss": "3.3.2",
-    "web-vitals": "^3.4.0"
+    "web-vitals": "^3.4.0",
+    "xz-decompress": "^0.2.1"
   },
   "devDependencies": {
     "@tailwindcss/typography": "^0.5.9",

diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
diff --git a/src/utils/manifest.js b/src/utils/manifest.js
@@ -45,23 +45,21 @@ export class Image {
 
   constructor(json) {
     this.name = json.name
-    this.size = json.size
     this.sparse = json.sparse
 
     if (this.name === 'system') {
       this.checksum = json.alt.hash
       this.fileName = `${this.name}-skip-chunks-${json.hash_raw}.img`
+      this.archiveUrl = json.alt.url
+      this.size = json.alt.size
     } else {
       this.checksum = json.hash
       this.fileName = `${this.name}-${json.hash_raw}.img`
-    }
+      this.archiveUrl = json.url
+      this.size = json.size
+    } 
 
-    let baseUrl = json.url.split('/')
-    baseUrl.pop()
-    baseUrl = baseUrl.join('/')
-
-    this.archiveFileName = this.fileName + '.gz'
-    this.archiveUrl = `${baseUrl}/${this.archiveFileName}`
+    this.archiveFileName = this.archiveUrl.split('/').pop()
   }
 }
 

diff --git a/src/utils/manifest.test.js b/src/utils/manifest.test.js
@@ -1,23 +1,54 @@
-import { expect, test } from 'vitest'
+import { expect, test, vi } from 'vitest'
 
-import jsSHA from 'jssha'
-import pako from 'pako'
+import * as Comlink from 'comlink'
 
 import config from '../config'
 import { getManifest } from './manifest'
 
+async function getImageWorker() {
+  let imageWorker
+
+  vi.mock('comlink')
+  vi.mocked(Comlink.expose).mockImplementation(worker => {
+    imageWorker = worker
+    imageWorker.init()
+  })
+
+  await import('./../workers/image.worker')
+
+  return imageWorker
+}
+
 for (const [branch, manifestUrl] of Object.entries(config.manifests)) {
   describe(`${branch} manifest`, async () => {
+    const imageWorkerFileHandler = {
+      getFile: vi.fn(),
+      createWritable: vi.fn().mockImplementation(() => ({
+        write: vi.fn(),
+        close: vi.fn(),
+      })),
+    }
+
+    global.navigator = {
+      storage: {
+        getDirectory: () => ({
+          getFileHandle: () => imageWorkerFileHandler,
+        })
+      }
+    }
+
+    const imageWorker = await getImageWorker()
+
     const images = await getManifest(manifestUrl)
 
     // Check all images are present
     expect(images.length).toBe(7)
 
     for (const image of images) {
       describe(`${image.name} image`, async () => {
-        test('gzip archive', () => {
-          expect(image.archiveFileName, 'archive to be a gzip').toContain('.gz')
-          expect(image.archiveUrl, 'archive url to be a gzip').toContain('.gz')
+        test('xz archive', () => {
+          expect(image.archiveFileName, 'archive to be in xz format').toContain('.xz')
+          expect(image.archiveUrl, 'archive url to be in xz format').toContain('.xz')
         })
 
         if (image.name === 'system') {
@@ -29,29 +60,14 @@ for (const [branch, manifestUrl] of Object.entries(config.manifests)) {
         }
 
         test('image and checksum', async () => {
-          const response = await fetch(image.archiveUrl)
-          expect(response.ok, 'to be uploaded').toBe(true)
-
-          const inflator = new pako.Inflate()
-          const shaObj = new jsSHA('SHA-256', 'UINT8ARRAY')
+          imageWorkerFileHandler.getFile.mockImplementation(async () => {
+            const response = await fetch(image.archiveUrl)
+            expect(response.ok, 'to be uploaded').toBe(true)
 
-          inflator.onData = function (chunk) {
-            shaObj.update(chunk)
-          }
-
-          inflator.onEnd = function (status) {
-            expect(status, 'to decompress').toBe(0)
-
-            const checksum = shaObj.getHash('HEX')
-            expect(checksum, 'to match').toBe(image.checksum)
-          }
-
-          const reader = response.body.getReader()
-          while (true) {
-            const { done, value } = await reader.read()
-            if (done) break
-            inflator.push(value)
-          }
+            return response.blob()
+          })
+
+          await imageWorker.unpackImage(image)
         }, 8 * 60 * 1000)
       })
     }

diff --git a/src/workers/image.worker.js b/src/workers/image.worker.js
@@ -1,7 +1,7 @@
 import * as Comlink from 'comlink'
 
 import jsSHA from 'jssha'
-import pako from 'pako'
+import { XzReadableStream } from 'xz-decompress';
 
 import { Image } from '@/utils/manifest'
 
@@ -10,7 +10,7 @@ import { Image } from '@/utils/manifest'
  *
  * @callback chunkCallback
  * @param {Uint8Array} chunk
- * @returns {void}
+ * @returns {Promise<void>}
  */
 
 /**
@@ -35,7 +35,7 @@ async function readChunks(reader, total, { onChunk, onProgress = undefined }) {
   while (true) {
     const { done, value } = await reader.read()
     if (done) break
-    onChunk(value)
+    await onChunk(value)
     processed += value.length
     onProgress?.(processed / total)
   }
@@ -83,7 +83,7 @@ const imageWorker = {
       const contentLength = +response.headers.get('Content-Length')
       const reader = response.body.getReader()
       await readChunks(reader, contentLength, {
-        onChunk: (chunk) => writable.write(chunk),
+        onChunk: async (chunk) => await writable.write(chunk),
         onProgress,
       })
       onProgress?.(1)
@@ -108,7 +108,7 @@ const imageWorker = {
    * @returns {Promise<void>}
    */
   async unpackImage(image, onProgress = undefined) {
-    const { archiveFileName, checksum: expectedChecksum, fileName } = image
+    const { archiveFileName, checksum: expectedChecksum, fileName, size: imageSize } = image
 
     let archiveFile
     try {
@@ -129,28 +129,18 @@ const imageWorker = {
     const shaObj = new jsSHA('SHA-256', 'UINT8ARRAY')
     let complete
     try {
-      const reader = archiveFile.stream().getReader()
-      await new Promise(async (resolve, reject) => {
-        const inflator = new pako.Inflate()
-        inflator.onData = function (chunk) {
-          writable.write(chunk)
+      const reader = (new XzReadableStream(archiveFile.stream())).getReader()
+
+      await readChunks(reader, imageSize, {
+        onChunk: async (chunk) => {
+          await writable.write(chunk)
           shaObj.update(chunk)
-        }
-        inflator.onEnd = function (status) {
-          if (status) {
-            reject(`Decompression error ${status}: ${inflator.msg}`)
-          } else {
-            resolve()
-          }
-          complete = true
-        }
-
-        await readChunks(reader, archiveFile.size, {
-          onChunk: (chunk) => inflator.push(chunk),
-          onProgress,
-        })
-        onProgress?.(1)
+        },
+        onProgress,
       })
+
+      complete = true
+      onProgress?.(1)
     } catch (e) {
       throw `Error unpacking archive: ${e}`
     }