Browse Source

Detect invalid content (#458)

Lutz Roeder 3 years ago
parent
commit
eb1588e76f
3 changed files with 20 additions and 2 deletions
  1. 10 1
      source/onnx.js
  2. 3 1
      source/view.js
  3. 7 0
      test/models.json

+ 10 - 1
source/onnx.js

@@ -154,6 +154,12 @@ onnx.ModelFactory = class {
                 }
             }
         }
+        if (extension === 'onnx' && stream.length > 3) {
+            const signature = stream.peek(2);
+            if (signature[0] === 0x80 && signature[1] < 7) {
+                return 'onnx.pickle';
+            }
+        }
         return undefined;
     }
 
@@ -273,8 +279,11 @@ onnx.ModelFactory = class {
                     }
                 });
             }
+            case 'onnx.pickle': {
+                return Promise.reject(new onnx.Error('Unsupported Pickle content.'));
+            }
             default: {
-                throw new onnx.Error("Unsupported ONNX format '" + match + "'.");
+                return Promise.reject(new onnx.Error("Unsupported ONNX format '" + match + "'."));
             }
         }
     }

+ 3 - 1
source/view.js

@@ -1807,6 +1807,7 @@ view.ModelFactoryService = class {
                 const formats = [
                     { name: 'onnxruntime.experimental.fbs.InferenceSession data', identifier: 'ORTM' },
                     { name: 'tflite.Model data', identifier: 'TFL3' },
+                    { name: 'torch.jit.mobile.serialization.Module data', identifier: 'PTMF' }, // https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/serialization/mobile_bytecode.fbs
                     { name: 'FlatBuffers ENNC data', identifier: 'ENNC' },
                 ];
                 for (const format of formats) {
@@ -2084,7 +2085,8 @@ view.ModelFactoryService = class {
                 { name: 'TSD header', value: /^%TSD-Header-###%/ },
                 { name: 'AppleDouble data', value: /^\x00\x05\x16\x07/ },
                 { name: 'TensorFlow Hub module', value: /^\x08\x03$/, identifier: 'tfhub_module.pb' },
-                { name: 'ViSQOL model', value: /^svm_type\snu_svr/ }
+                { name: 'ViSQOL model', value: /^svm_type\snu_svr/ },
+                { name: 'SenseTime model', value: /^STEF/ }
             ];
             /* eslint-enable no-control-regex */
             const buffer = stream.peek(Math.min(4096, stream.length));

+ 7 - 0
test/models.json

@@ -1,4 +1,11 @@
 [
+  {
+    "type":     "_",
+    "target":   "boolean.onnx",
+    "source":   "https://github.com/lutzroeder/netron/files/8985631/boolean.onnx.zip[boolean.onnx]",
+    "error":    "Unsupported Pickle content in 'boolean.onnx'.",
+    "link":     "https://github.com/lutzroeder/netron/issues/458"
+  },
   {
     "type":     "_",
     "target":   "config.json",