Support Protocol5 (#4)

* support protocol 5

* PObject enhancement

* update docs

* fix typo

* add more ut

* update docs

README.md

@@ -9,7 +9,7 @@ A pure Javascript implemented parser for [Python pickle format](https://docs.pyt
 
 ## Features
 
-* Fully supports [Pickle protocol version 4](https://peps.python.org/pep-3154/) opcodes.
+* Fully supports Pickle protocol version 0~5 opcodes.
 * Pure Typescript implemented.
 * Provides `ParserOptions` to customize Unpickling.
 * Supports Browser.
@@ -23,6 +23,7 @@ A pure Javascript implemented parser for [Python pickle format](https://docs.pyt
 * [Pickle protocol version 2 (Python 2.3)](https://peps.python.org/pep-0307/)
 * Pickle protocol version 3 (Python 3.0)
 * [Pickle protocol version 4 (Python 3.4)](https://peps.python.org/pep-3154/)
+* [Pickle protocol version 5 (Python 3.8)](https://peps.python.org/pep-0574/)
 
 For more details, see: [Supported Opcodes](./SUPPORTED_OPCODES.md)
 

SUPPORTED_OPCODES.md

@@ -67,6 +67,6 @@
 |  STACK_GLOBAL       |  \x93   |  147   |  ✅        |  Protocol 4   |
 |  MEMOIZE            |  \x94   |  148   |  ✅        |  Protocol 4   |
 |  FRAME              |  \x95   |  149   |  ✅        |  Protocol 4   |
-|  BYTEARRAY8         |  \x96   |  150   |  ❌        |  Protocol 5   |
-|  NEXT_BUFFER        |  \x97   |  151   |  ❌        |  Protocol 5   |
-|  READONLY_BUFFER    |  \x98   |  152   |  ❌        |  Protocol 5   |
+|  BYTEARRAY8         |  \x96   |  150   |  ✅        |  Protocol 5   |
+|  NEXT_BUFFER        |  \x97   |  151   |  ✅        |  Protocol 5   |
+|  READONLY_BUFFER    |  \x98   |  152   |  ✅        |  Protocol 5   |

examples/testprotocol5.py

@@ -0,0 +1,20 @@
+import pickle
+
+class mybytearray(bytearray):
+
+   def __reduce_ex__(self, protocol):
+      if protocol >= 5:
+         return type(self), (pickle.PickleBuffer(self),), None
+
+key = mybytearray([0x01, 0x02, 0x03, 0x04])
+filehandler = open(b"bytearray8.pkl", "wb")
+
+buffers = []
+t= memoryview(bytearray())
+oo = [pickle.PickleBuffer(memoryview(bytearray()).toreadonly()), pickle.PickleBuffer(bytearray())]
+d = pickle.dumps(oo, pickle.HIGHEST_PROTOCOL, buffer_callback= lambda _: False)
+pickle.dump(mybytearray([0x01, 0x02, 0x03, 0x04]), filehandler, pickle.HIGHEST_PROTOCOL, buffer_callback= lambda _: True)
+print(buffers)
+buffers = [bytearray([]), []]
+l = pickle.loads(d, buffers=buffers)
+print(l)

src/PObject.ts

@@ -1,5 +1,8 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
-export function createPObject(module: string, name: string): (new (...args: any[]) => any) | ((...args: any[]) => any) {
+export function createPObject<T extends (new (...args: any[]) => any) | ((...args: any[]) => any)>(
+    module: string,
+    name: string,
+): T {
     const PObject = function (this: any, ...args: any[]): any {
         if (new.target) {
             Object.defineProperty(this, 'args', {
@@ -21,10 +24,10 @@ export function createPObject(module: string, name: string): (new (...args: any[
             PFunction.prototype.__name__ = name;
             return Reflect.construct(PFunction, args);
         }
-    } as unknown as (new (...args: any[]) => any) | ((...args: any[]) => any);
+    } as T;
     PObject.prototype.__module__ = module;
     PObject.prototype.__name__ = name;
-    PObject.prototype.__setnewargs_ex__ = function (kwargs: any) {
+    PObject.prototype.__setnewargs_ex__ = function (...kwargs: any) {
         Object.defineProperty(this, 'kwargs', {
             value: kwargs,
             enumerable: false,

src/parser.ts

@@ -24,6 +24,7 @@ export interface ParserOptions {
     extensionResolver: ExtensionResolver;
     unpicklingTypeOfSet: UnpicklingTypeOfSet;
     unpicklingTypeOfDictionary: UnpicklingTypeOfDictionary;
+    buffers?: Iterator<any>;
 }
 
 const DefualtOptions: ParserOptions = {
@@ -51,6 +52,7 @@ export class Parser {
     private readonly _extensionResolver: ExtensionResolver;
     private readonly _setProvider: ISetProvider;
     private readonly _dictionaryProvider: IDictionaryProvider;
+    private readonly _buffers?: Iterator<any>;
 
     constructor(options?: Partial<ParserOptions>) {
         this._options = { ...DefualtOptions, ...options };
@@ -59,6 +61,7 @@ export class Parser {
         this._extensionResolver = this._options.extensionResolver;
         this._setProvider = SetProviderFactory(this._options.unpicklingTypeOfSet);
         this._dictionaryProvider = DictionaryProviderFactory(this._options.unpicklingTypeOfDictionary);
+        this._buffers = options?.buffers;
     }
 
     parse<T>(buffer: Uint8Array | Int8Array | Uint8ClampedArray): T {
@@ -72,14 +75,14 @@ export class Parser {
         const memo = new Map();
         while (reader.hasNext()) {
             const opcode = reader.byte();
-            // console.log(`${(reader.position - 1).toString()} ${opcode}`);
+            // console.log(`${((reader as any)._position - 1).toString()} ${opcode}`);
             // console.log('metastack:', metastack, '\nstack:', stack);
             // console.log('\nmemo:', Array.from(memo.entries()));
             switch (opcode) {
                 // Structural
                 case OP.PROTO: {
                     const version = reader.byte();
-                    if (version > 4) {
+                    if (version > 5) {
                         throw new Error(`Unsupported protocol version '${version}'.`);
                     }
                     break;
@@ -462,6 +465,26 @@ export class Parser {
                     break;
                 }
 
+                case OP.BYTEARRAY8:
+                    stack.push(reader.bytes(reader.uint64()));
+                    break;
+
+                case OP.NEXT_BUFFER: {
+                    if (this._buffers == null) {
+                        throw new Error('pickle stream refers to out-of-band data but no *buffers* argument was given');
+                    }
+                    const next = this._buffers.next();
+                    if (next.done) {
+                        throw new Error('not enough out-of-band buffers');
+                    }
+                    stack.push(next.value);
+                    break;
+                }
+
+                case OP.READONLY_BUFFER:
+                    stack.push(stack.pop());
+                    break;
+
                 default:
                     throw new Error(`Unsupported opcode '${opcode}'.`);
             }

test/PObject.test.ts

@@ -0,0 +1,50 @@
+import { createPObject } from '../src/PObject';
+
+describe('PObject', () => {
+    describe('createPObject', () => {
+        it('can be created', () => {
+            const name = 'name';
+            const module = 'module';
+            const pobject = createPObject(module, name);
+            expect(pobject).toBeDefined();
+        });
+
+        it('can be used with as a class', () => {
+            const name = 'name';
+            const module = 'module';
+            const data = [1, true, [null, 'str']];
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            const pobject = createPObject<new (...args: any[]) => any>(module, name);
+            const obj = new pobject(...data);
+            expect(obj).toHaveProperty('__module__', module);
+            expect(obj).toHaveProperty('__name__', name);
+            expect(obj.args).toStrictEqual(data);
+        });
+
+        it('can be used with as a function', () => {
+            const name = 'name';
+            const module = 'module';
+            const data = [1, true, [null, 'str']];
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            const pobject = createPObject<(...args: any[]) => any>(module, name);
+            const obj = pobject(...data);
+            expect(obj).toHaveProperty('__module__', module);
+            expect(obj).toHaveProperty('__name__', name);
+            expect(obj.args).toStrictEqual(data);
+        });
+
+        it('can be worked with __setnewargs_ex__', () => {
+            const name = 'name';
+            const module = 'module';
+            const data = [1, true, [null, 'str']];
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            const pobject = createPObject<new (...args: any[]) => any>(module, name);
+            const obj = new pobject(...data);
+            obj.__setnewargs_ex__(...data);
+            expect(obj).toHaveProperty('__module__', module);
+            expect(obj).toHaveProperty('__name__', name);
+            expect(obj.args).toStrictEqual(data);
+            expect(obj.kwargs).toStrictEqual(data);
+        });
+    });
+});

test/integration/_runner.py

@@ -3,5 +3,5 @@
 
 if __name__ == "__main__":
     func = getattr(__import__(sys.argv[1]), sys.argv[2])
-    sys.stdout.buffer.write(pickle.dumps(func()))
+    sys.stdout.buffer.write(pickle.dumps(func(), pickle.HIGHEST_PROTOCOL, buffer_callback= lambda _: False))
     sys.stdout.flush()

test/integration/intergration.test.ts

@@ -85,3 +85,76 @@ describe('long', () => {
         expect(obj.toString()).toStrictEqual(expected);
     });
 });
+
+describe('protocol5', () => {
+    it('correctly unpickl bytearray8', async () => {
+        const expected = Buffer.from([1, 2, 3]);
+        const data = await caller('protocol5', 'bytearray8');
+        const obj = new Parser().parse<Buffer>(data);
+        expect(obj).toStrictEqual(expected);
+    });
+
+    it('correctly unpickl next_buffer', async () => {
+        const expected = 123;
+        const data = await caller('protocol5', 'next_buffer');
+        const obj = new Parser({
+            buffers: (function* () {
+                yield expected;
+            })(),
+        }).parse(data);
+        expect(obj).toStrictEqual(expected);
+    });
+
+    it('correctly unpickl multi_next_buffer', async () => {
+        const expected = [123, 'str'];
+        const data = await caller('protocol5', 'multi_next_buffer');
+        const obj = new Parser({
+            buffers: expected.values(),
+        }).parse(data);
+        expect(obj).toStrictEqual(expected);
+    });
+
+    it('correctly unpickl readonly_buffer', async () => {
+        const expected = 123;
+        const data = await caller('protocol5', 'readonly_buffer');
+        const obj = new Parser({
+            buffers: (function* () {
+                yield expected;
+            })(),
+        }).parse(data);
+        expect(obj).toStrictEqual(expected);
+    });
+
+    it('correctly unpickl next_buffer_and_readonly_buffer', async () => {
+        const expected = [123, [1, '22', null]];
+        const data = await caller('protocol5', 'next_buffer_and_readonly_buffer');
+        const obj = new Parser({
+            buffers: expected.values(),
+        }).parse(data);
+        expect(obj).toStrictEqual(expected);
+    });
+
+    it('correctly unpickl next_buffer_with_reduce_ex', async () => {
+        class mybytearray {
+            public args: number[];
+            constructor(args: number[]) {
+                this.args = args;
+            }
+            static __reduce_ex__(args: number[]) {
+                return new mybytearray(args);
+            }
+        }
+        const externalData = [1, 2, 3, 4];
+        const expected = new mybytearray(externalData);
+        const registry = new NameRegistry();
+        registry.register('protocol5', 'mybytearray', mybytearray.__reduce_ex__);
+        const data = await caller('protocol5', 'next_buffer_with_reduce_ex');
+        const obj = new Parser({
+            nameResolver: registry,
+            buffers: (function* () {
+                yield externalData;
+            })(),
+        }).parse<mybytearray>(data);
+        expect(obj).toStrictEqual(expected);
+    });
+});

test/integration/protocol5.py

@@ -0,0 +1,30 @@
+import pickle
+
+def bytearray8():
+    key = bytearray([0x01, 0x02, 0x03])
+    return key
+
+def next_buffer():
+    key = pickle.PickleBuffer(bytearray())
+    return key
+
+def multi_next_buffer():
+    key = [pickle.PickleBuffer(bytearray()), pickle.PickleBuffer(bytearray())]
+    return key
+
+def readonly_buffer():
+    key = pickle.PickleBuffer(memoryview(bytearray()).toreadonly())
+    return key
+
+def next_buffer_and_readonly_buffer():
+    key = [pickle.PickleBuffer(memoryview(bytearray()).toreadonly()), pickle.PickleBuffer(bytearray())]
+    return key
+
+class mybytearray(bytearray):
+    def __reduce_ex__(self, protocol):
+        if protocol >= 5:
+            return type(self), (pickle.PickleBuffer(self),), None
+
+def next_buffer_with_reduce_ex():
+    key = mybytearray([0x01, 0x02, 0x03, 0x04])
+    return key

test/parser.test.ts

@@ -95,5 +95,34 @@ describe('Parser', () => {
             const pkl = new Uint8Array([OP.PROTO, 4, OP.BININT1, 3, OP.STOP]);
             expect(parser.parse(pkl)).toEqual(3);
         });
+
+        it('correctly load buffers', () => {
+            const expected = [123, 'str'];
+            const parser = new Parser({
+                buffers: (function* () {
+                    yield expected;
+                })(),
+            });
+            const pkl = new Uint8Array([OP.PROTO, 5, OP.NEXT_BUFFER, OP.STOP]);
+            expect(parser.parse(pkl)).toStrictEqual(expected);
+        });
+
+        it('throws errors if buffers not provided', () => {
+            const parser = new Parser();
+            const pkl = new Uint8Array([OP.PROTO, 5, OP.NEXT_BUFFER, OP.STOP]);
+            expect(() => parser.parse(pkl)).toThrow(
+                'pickle stream refers to out-of-band data but no *buffers* argument was given',
+            );
+        });
+
+        it('throws errors if not enough buffers', () => {
+            const parser = new Parser({
+                buffers: (function* () {
+                    yield 1;
+                })(),
+            });
+            const pkl = new Uint8Array([OP.PROTO, 5, OP.NEXT_BUFFER, OP.NEXT_BUFFER, OP.STOP]);
+            expect(() => parser.parse(pkl)).toThrow('not enough out-of-band buffers');
+        });
     });
 });