deser_pickle/
ser.rs

1use alloc::collections::BTreeMap;
2use alloc::format;
3use alloc::string::String;
4use alloc::vec::Vec;
5
6use deser_core::State;
7use deser_core::ext::{BigInt, ExtValue, Number};
8use deser_core::ser::{self, EventSink, SerializeDriver, SerializeRef};
9use deser_core::{Atom, Error, ErrorKind, Event, Serialize};
10
11use crate::compat;
12use crate::types::{ClassData, Form, FormData, Global, Kind, KindData, Reference, SharedIdData};
13use crate::vm::HIGHEST_PROTOCOL;
14
15/// The protocol written by default.
16const DEFAULT_PROTOCOL: u8 = 4;
17
18/// Configures how values are serialized.
19///
20/// The output can be read with Python's `pickle.loads` (see the
21/// [crate documentation](crate#serialization)).  The options are the
22/// [protocol](Self::set_protocol) and the [`Context`](deser_core::Context).
23///
24/// ```
25/// use deser_pickle::SerializerConfig;
26///
27/// const CONFIG: SerializerConfig = SerializerConfig::builder().protocol(2).build();
28/// assert_eq!(CONFIG.to_vec(&vec![1, 2]).unwrap(), b"\x80\x02](K\x01K\x02e.");
29/// ```
30#[derive(Debug, Clone, PartialEq, Eq)]
31pub struct SerializerConfig {
32    context: deser_core::Context,
33    protocol: u8,
34}
35
36impl Default for SerializerConfig {
37    fn default() -> SerializerConfig {
38        SerializerConfig::new()
39    }
40}
41
42impl SerializerConfig {
43    /// Creates the default configuration.
44    pub const fn new() -> SerializerConfig {
45        SerializerConfig {
46            context: deser_core::Context::new(),
47            protocol: DEFAULT_PROTOCOL,
48        }
49    }
50
51    /// Returns a builder for the configuration (see [`SerializerConfigBuilder`]).
52    pub const fn builder() -> SerializerConfigBuilder {
53        SerializerConfigBuilder::new()
54    }
55
56    /// Returns a builder that starts with this configuration.
57    pub const fn into_builder(self) -> SerializerConfigBuilder {
58        SerializerConfigBuilder { value: self }
59    }
60
61    /// Sets the context the values are serialized in.
62    ///
63    /// The values of the context are the defaults of the extension values
64    /// of the state (see [`Context`](deser_core::Context)).  The
65    /// serializers and writers created with the configuration use this
66    /// context.  A context set on the driver takes precedence.
67    pub fn set_context(&mut self, context: deser_core::Context) {
68        self.context = context;
69    }
70
71    /// Returns the context the values are serialized in.
72    pub fn context(&self) -> &deser_core::Context {
73        &self.context
74    }
75
76    /// Sets the protocol that is written.
77    ///
78    /// Protocols 2 to 5 are supported, the default is 4 (which Python
79    /// reads since 3.4).  Protocol 2 can be read by Python 2.  Serializing
80    /// fails with other protocols.
81    pub fn set_protocol(&mut self, protocol: u8) {
82        self.protocol = protocol;
83    }
84
85    /// Returns the protocol that is written.
86    pub fn protocol(&self) -> u8 {
87        self.protocol
88    }
89
90    /// Gives the context to a driver which has none.
91    #[inline]
92    fn apply_context(&self, driver: &mut SerializeDriver<'_>) {
93        if !self.context.is_empty() {
94            driver.set_default_context(self.context.clone());
95        }
96    }
97
98    /// Serializes the given value.
99    pub fn to_vec<T: Serialize + ?Sized>(&self, value: &T) -> Result<Vec<u8>, Error> {
100        self.to_vec_ref(SerializeRef::new(&value))
101    }
102
103    /// Serializes the given value with a configured driver.
104    ///
105    /// The callback is invoked with the driver before the serialization
106    /// starts, for instance to add [`Layer`](deser_core::ser::Layer)s.
107    pub fn to_vec_with<F, T: Serialize + ?Sized>(
108        &self,
109        value: &T,
110        setup: F,
111    ) -> Result<Vec<u8>, Error>
112    where
113        F: FnOnce(&mut SerializeDriver<'_>),
114    {
115        let mut driver = SerializeDriver::new(&value);
116        setup(&mut driver);
117        self.apply_context(&mut driver);
118        serialize_driver(&mut driver, self.protocol)
119    }
120
121    /// Serializes a value whose type is erased (see
122    /// [`to_vec`](Self::to_vec)).
123    fn to_vec_ref(&self, value: SerializeRef<'_>) -> Result<Vec<u8>, Error> {
124        let mut driver = SerializeDriver::from_ref(value);
125        self.apply_context(&mut driver);
126        serialize_driver(&mut driver, self.protocol)
127    }
128}
129
130/// Builds a [`SerializerConfig`].
131///
132/// The methods have the names of the setters of [`SerializerConfig`] (without `set_`).
133#[derive(Debug, Clone)]
134#[must_use]
135pub struct SerializerConfigBuilder {
136    value: SerializerConfig,
137}
138
139impl SerializerConfigBuilder {
140    /// Creates a builder that starts with the default.
141    pub const fn new() -> SerializerConfigBuilder {
142        SerializerConfigBuilder {
143            value: SerializerConfig::new(),
144        }
145    }
146
147    /// Sets the context the values are serialized in.
148    ///
149    /// See [`SerializerConfig::set_context`].
150    pub fn context(mut self, context: deser_core::Context) -> SerializerConfigBuilder {
151        self.value.set_context(context);
152        self
153    }
154
155    /// Sets the protocol that is written.
156    ///
157    /// See [`SerializerConfig::set_protocol`].
158    pub const fn protocol(mut self, protocol: u8) -> SerializerConfigBuilder {
159        self.value.protocol = protocol;
160        self
161    }
162
163    /// Returns the built [`SerializerConfig`].
164    pub const fn build(self) -> SerializerConfig {
165        // the value cannot be moved out of the builder in a const fn as the
166        // builder needs dropping (the context has a destructor)
167        // SAFETY: the value is read once and the builder is forgotten
168        let value = unsafe { core::ptr::read(&self.value) };
169        core::mem::forget(self);
170        value
171    }
172}
173
174impl Default for SerializerConfigBuilder {
175    fn default() -> SerializerConfigBuilder {
176        SerializerConfigBuilder::new()
177    }
178}
179
180/// Serializes values into pickles.
181///
182/// Every value is a pickle of its own, they are appended to the output one
183/// after another (they can be read back with
184/// [`Deserializer::iter`](crate::Deserializer::iter) or by calling
185/// Python's `pickle.load` repeatedly).
186///
187/// ```
188/// use deser_pickle::Serializer;
189///
190/// let mut serializer = Serializer::new();
191/// serializer.serialize(&true).unwrap();
192/// serializer.serialize(&"hi").unwrap();
193/// assert_eq!(serializer.output(), b"\x80\x04\x88.\x80\x04\x8c\x02hi.");
194/// ```
195///
196/// The serializer is also the stream serializer of the format (see
197/// [`StreamSerializer`](ser::StreamSerializer)).  To write to a
198/// [`Write`](std::io::Write) use [`SerializerConfig::writer`].
199#[derive(Debug, Clone, Default)]
200pub struct Serializer {
201    config: SerializerConfig,
202    out: Vec<u8>,
203}
204
205impl Serializer {
206    /// Creates a serializer.
207    pub fn new() -> Serializer {
208        Serializer::with_config(SerializerConfig::new())
209    }
210
211    /// Creates a serializer with the given configuration.
212    pub fn with_config(config: SerializerConfig) -> Serializer {
213        Serializer {
214            config,
215            out: Vec::new(),
216        }
217    }
218
219    /// Returns the configuration.
220    pub fn config(&self) -> &SerializerConfig {
221        &self.config
222    }
223
224    /// Serializes a value and appends it to the output.
225    pub fn serialize<T: Serialize + ?Sized>(&mut self, value: &T) -> Result<(), Error> {
226        ser::Serializer::serialize(self, value)
227    }
228
229    /// Serializes a value with a configured driver.
230    ///
231    /// The callback is invoked with the driver before the value is
232    /// serialized, for instance to add [`Layer`](deser_core::ser::Layer)s.
233    pub fn serialize_with<F, T: Serialize + ?Sized>(
234        &mut self,
235        value: &T,
236        setup: F,
237    ) -> Result<(), Error>
238    where
239        F: FnOnce(&mut SerializeDriver<'_>),
240    {
241        ser::Serializer::serialize_with(self, value, setup)
242    }
243
244    /// Returns the output written so far (that was not cleared).
245    pub fn output(&self) -> &[u8] {
246        &self.out
247    }
248
249    /// Returns the output.
250    pub fn finish(self) -> Vec<u8> {
251        self.out
252    }
253}
254
255impl ser::Serializer for Serializer {
256    fn drive(&mut self, driver: &mut SerializeDriver<'_>) -> Result<(), Error> {
257        self.config.apply_context(driver);
258        let bytes = serialize_driver(driver, self.config.protocol)?;
259        self.out.extend_from_slice(&bytes);
260        Ok(())
261    }
262}
263
264impl ser::StreamSerializer for Serializer {
265    fn output(&self) -> &[u8] {
266        &self.out
267    }
268
269    fn clear_output(&mut self) {
270        self.out.clear();
271    }
272}
273
274#[cfg(feature = "io")]
275impl SerializerConfig {
276    /// Creates a writer of values (see
277    /// [`deser::io::Writer`](deser_core::io::Writer)).
278    ///
279    /// Every value is written as a pickle once it's complete.
280    pub fn writer<W: std::io::Write>(&self, writer: W) -> deser_core::io::Writer<W, Serializer> {
281        deser_core::io::Writer::new(writer, Serializer::with_config(self.clone()))
282    }
283
284    /// Serializes a value to a writer.
285    ///
286    /// See [`to_writer`].
287    pub fn to_writer<W: std::io::Write, T: Serialize + ?Sized>(
288        &self,
289        writer: W,
290        value: &T,
291    ) -> Result<(), Error> {
292        self.writer(writer).write(value)
293    }
294}
295
296/// Serializes a value to a writer.
297///
298/// ```
299/// let mut out = Vec::new();
300/// deser_pickle::to_writer(&mut out, &vec!["a", "b"]).unwrap();
301/// assert_eq!(out, b"\x80\x04](\x8c\x01a\x8c\x01be.");
302/// ```
303#[cfg(feature = "io")]
304pub fn to_writer<W: std::io::Write, T: Serialize + ?Sized>(
305    writer: W,
306    value: &T,
307) -> Result<(), Error> {
308    SerializerConfig::new().to_writer(writer, value)
309}
310
311/// Serializes a value into a pickle.
312///
313/// This uses the default [`SerializerConfig`] (protocol 4).
314///
315/// ```
316/// #[derive(deser::Serialize)]
317/// struct Package {
318///     name: String,
319///     tags: Vec<String>,
320/// }
321/// let package = Package { name: "deser".into(), tags: vec!["a".into()] };
322/// assert_eq!(
323///     deser_pickle::to_vec(&package).unwrap(),
324///     b"\x80\x04}(\x8c\x04name\x8c\x05deser\x8c\x04tags](\x8c\x01aeu."
325/// );
326/// ```
327pub fn to_vec<T: Serialize + ?Sized>(value: &T) -> Result<Vec<u8>, Error> {
328    SerializerConfig::new().to_vec(value)
329}
330
331/// Serializes the value of a driver.
332fn serialize_driver(driver: &mut SerializeDriver<'_>, protocol: u8) -> Result<Vec<u8>, Error> {
333    if !(2..=HIGHEST_PROTOCOL).contains(&protocol) {
334        return Err(Error::new(
335            ErrorKind::Configuration,
336            format!("unsupported pickle protocol {}", protocol),
337        ));
338    }
339    let mut writer = Writer {
340        out: alloc::vec![0x80, protocol],
341        proto: protocol,
342        stack: Vec::new(),
343        memo: BTreeMap::new(),
344        memo_len: 0,
345        skip: 0,
346        started: false,
347    };
348    driver.drive_sink(&mut writer)?;
349    writer.finish()
350}
351
352/// What is written when a container ends.
353#[derive(Clone, Copy, PartialEq)]
354enum End {
355    /// A list or the items of a list-like object (`APPENDS`).
356    Appends,
357    /// A dict or the items of a dict-like object (`SETITEMS`).
358    SetItems,
359    /// A set (`ADDITEMS`).
360    AddItems,
361    /// A tuple (`TUPLE`).
362    Tuple,
363    /// A frozenset (`FROZENSET`).
364    FrozenSet,
365    /// The list of a set or frozenset of protocols 2 and 3 (`LIST`).
366    List,
367    /// A map that must be empty (keyword arguments before protocol 4).
368    Empty,
369}
370
371/// The opcodes written after the end of a container.
372#[derive(Clone, Copy, Default)]
373struct Suffix {
374    buf: [u8; 4],
375    len: u8,
376}
377
378impl Suffix {
379    fn new(bytes: &[u8]) -> Suffix {
380        Suffix::default().with(bytes)
381    }
382
383    fn with(mut self, bytes: &[u8]) -> Suffix {
384        let len = usize::from(self.len);
385        self.buf[len..len + bytes.len()].copy_from_slice(bytes);
386        self.len += bytes.len() as u8;
387        self
388    }
389
390    fn as_slice(&self) -> &[u8] {
391        &self.buf[..usize::from(self.len)]
392    }
393}
394
395/// An open container.
396struct Open {
397    end: End,
398    suffix: Suffix,
399    /// The id to memoize the value with after the suffix (for the values
400    /// that are only complete at their end).
401    memoize: Option<u64>,
402    is_map: bool,
403    /// `true` if a key comes next (maps only).
404    expects_key: bool,
405    /// `true` if the value is (in) a key, which must be hashable.
406    in_key: bool,
407}
408
409/// Writes the events of a value.
410struct Writer {
411    out: Vec<u8>,
412    proto: u8,
413    stack: Vec<Open>,
414    /// The memo index of the values with a shared id.
415    memo: BTreeMap<u64, u32>,
416    memo_len: u32,
417    /// The depth of a value that is skipped as it was written before.
418    skip: usize,
419    /// `true` once the top-level value started.
420    started: bool,
421}
422
423impl EventSink for Writer {
424    fn event(
425        &mut self,
426        event: Event<'_>,
427        _value: SerializeRef<'_>,
428        state: &mut State,
429    ) -> Result<(), Error> {
430        if self.skip > 0 {
431            match event {
432                Event::MapStart(_) | Event::SeqStart(_) => self.skip += 1,
433                Event::MapEnd | Event::SeqEnd => {
434                    self.skip -= 1;
435                    if self.skip == 0 {
436                        self.value_done();
437                    }
438                }
439                Event::Atom(_) => {}
440            }
441            return Ok(());
442        }
443        let Some(open) = self.stack.last() else {
444            if self.started {
445                return Err(Error::new(ErrorKind::InvalidState, "unexpected event"));
446            }
447            self.started = true;
448            return self.value(event, state, false);
449        };
450        let in_key = open.in_key || open.is_map && open.expects_key;
451        match event {
452            Event::MapEnd if open.is_map && open.expects_key => self.close(),
453            Event::SeqEnd if !open.is_map => self.close(),
454            Event::MapEnd | Event::SeqEnd => {
455                Err(Error::new(ErrorKind::InvalidState, "unexpected end event"))
456            }
457            _ if open.end == End::Empty => Err(Error::new(
458                ErrorKind::UnsupportedType,
459                "keyword arguments need protocol 4",
460            )),
461            event => self.value(event, state, in_key),
462        }
463    }
464}
465
466impl Writer {
467    fn finish(mut self) -> Result<Vec<u8>, Error> {
468        if !self.started || !self.stack.is_empty() || self.skip > 0 {
469            return Err(Error::new(ErrorKind::InvalidState, "incomplete value"));
470        }
471        self.out.push(b'.');
472        Ok(self.out)
473    }
474
475    /// Called when a value is complete.
476    fn value_done(&mut self) {
477        if let Some(open) = self.stack.last_mut()
478            && open.is_map
479        {
480            open.expects_key = !open.expects_key;
481        }
482    }
483
484    /// Writes the first event of a value.
485    fn value(&mut self, event: Event<'_>, state: &mut State, in_key: bool) -> Result<(), Error> {
486        let class = state.event::<ClassData>().and_then(|x| x.0.clone());
487        let form = state.event::<FormData>().and_then(|x| x.0);
488        let kind = state.event::<KindData>().and_then(|x| x.0);
489        let shared = state.event::<SharedIdData>().and_then(|x| x.0);
490        // values with an id are written once and referred to after that
491        if let Some(idx) = shared.and_then(|id| self.memo.get(&id).copied()) {
492            self.get(idx);
493            match event {
494                Event::MapStart(_) | Event::SeqStart(_) => self.skip = 1,
495                _ => self.value_done(),
496            }
497            return Ok(());
498        }
499        let Some(class) = class else {
500            return match event {
501                Event::Atom(atom) => {
502                    self.atom(atom, kind)?;
503                    self.value_done();
504                    Ok(())
505                }
506                Event::MapStart(_) => {
507                    if in_key {
508                        return Err(Error::new(
509                            ErrorKind::UnsupportedType,
510                            "maps cannot be keys of dicts",
511                        ));
512                    }
513                    self.out.push(b'}');
514                    self.memoize(shared);
515                    self.out.push(b'(');
516                    self.open(End::SetItems, Suffix::default(), None, true, false);
517                    Ok(())
518                }
519                Event::SeqStart(_) => self.open_seq(kind, in_key, shared, Suffix::default(), None),
520                Event::MapEnd | Event::SeqEnd => {
521                    Err(Error::new(ErrorKind::InvalidState, "unexpected end event"))
522                }
523            };
524        };
525        // objects are hashable (by their identity), their values do not
526        // need to be
527        self.global(&class)?;
528        let form = form.unwrap_or(match event {
529            Event::MapStart(_) => Form::State,
530            Event::SeqStart(_) => match kind {
531                None => Form::Items,
532                Some(Kind::Tuple) => Form::Arguments,
533                Some(_) => Form::Argument,
534            },
535            _ => Form::Argument,
536        });
537        match event {
538            Event::Atom(atom) => {
539                match form {
540                    Form::State | Form::Slots => {
541                        self.out.extend_from_slice(b")\x81");
542                        self.memoize(shared);
543                        self.atom(atom, kind)?;
544                        self.out.push(b'b');
545                    }
546                    _ => {
547                        self.atom(atom, kind)?;
548                        self.out.extend_from_slice(b"\x85R");
549                        self.memoize(shared);
550                    }
551                }
552                self.value_done();
553                Ok(())
554            }
555            Event::MapStart(_) => {
556                let (end, suffix, late) = match form {
557                    Form::State => {
558                        self.out.extend_from_slice(b")\x81");
559                        self.memoize(shared);
560                        self.out.push(b'}');
561                        (End::SetItems, Suffix::new(b"b"), None)
562                    }
563                    Form::Slots => {
564                        // the state is `(None, attributes)`
565                        self.out.extend_from_slice(b")\x81");
566                        self.memoize(shared);
567                        self.out.extend_from_slice(b"N}");
568                        (End::SetItems, Suffix::new(b"\x86b"), None)
569                    }
570                    Form::Items => {
571                        self.out.extend_from_slice(b")\x81");
572                        self.memoize(shared);
573                        (End::SetItems, Suffix::default(), None)
574                    }
575                    Form::Arguments if self.proto >= 4 => {
576                        self.out.extend_from_slice(b")}");
577                        (End::SetItems, Suffix::new(b"\x92"), shared)
578                    }
579                    Form::Arguments => {
580                        self.out.extend_from_slice(b")\x81");
581                        self.memoize(shared);
582                        (End::Empty, Suffix::default(), None)
583                    }
584                    Form::Argument => {
585                        self.out.push(b'}');
586                        (End::SetItems, Suffix::new(b"\x85R"), shared)
587                    }
588                };
589                if end != End::Empty {
590                    self.out.push(b'(');
591                }
592                self.open(end, suffix, late, true, false);
593                Ok(())
594            }
595            Event::SeqStart(_) => match form {
596                Form::State | Form::Slots => {
597                    self.out.extend_from_slice(b")\x81");
598                    self.memoize(shared);
599                    self.open_seq(kind, false, None, Suffix::new(b"b"), None)
600                }
601                Form::Items => {
602                    self.out.extend_from_slice(b")\x81");
603                    self.memoize(shared);
604                    self.out.push(b'(');
605                    let end = match kind {
606                        Some(Kind::Set | Kind::FrozenSet) => End::AddItems,
607                        _ => End::Appends,
608                    };
609                    self.open(end, Suffix::default(), None, false, false);
610                    Ok(())
611                }
612                Form::Arguments => {
613                    self.out.push(b'(');
614                    self.open(End::Tuple, Suffix::new(b"R"), shared, false, false);
615                    Ok(())
616                }
617                Form::Argument => self.open_seq(kind, false, None, Suffix::new(b"\x85R"), shared),
618            },
619            Event::MapEnd | Event::SeqEnd => {
620                Err(Error::new(ErrorKind::InvalidState, "unexpected end event"))
621            }
622        }
623    }
624
625    /// Opens a sequence of a kind.
626    ///
627    /// `memo` is the id of the sequence, `suffix` is written after it and
628    /// `late` is the id to memoize the value with after the suffix.
629    fn open_seq(
630        &mut self,
631        kind: Option<Kind>,
632        in_key: bool,
633        memo: Option<u64>,
634        suffix: Suffix,
635        late: Option<u64>,
636    ) -> Result<(), Error> {
637        let (end, suffix, late) = match (kind, in_key) {
638            (Some(Kind::Tuple), _) | (None, true) => (End::Tuple, suffix, memo.or(late)),
639            (Some(Kind::FrozenSet), _) | (Some(Kind::Set), true) => {
640                if self.proto >= 4 {
641                    (End::FrozenSet, suffix, memo.or(late))
642                } else {
643                    self.global(&Global::new(self.builtins(), "frozenset"))?;
644                    (
645                        End::List,
646                        Suffix::new(b"\x85R").with(suffix.as_slice()),
647                        memo.or(late),
648                    )
649                }
650            }
651            (Some(Kind::Set), false) => {
652                if self.proto >= 4 {
653                    self.out.push(0x8f);
654                    self.memoize(memo);
655                    (End::AddItems, suffix, late)
656                } else {
657                    self.global(&Global::new(self.builtins(), "set"))?;
658                    (
659                        End::List,
660                        Suffix::new(b"\x85R").with(suffix.as_slice()),
661                        memo.or(late),
662                    )
663                }
664            }
665            _ => {
666                self.out.push(b']');
667                self.memoize(memo);
668                (End::Appends, suffix, late)
669            }
670        };
671        self.out.push(b'(');
672        self.open(end, suffix, late, false, in_key);
673        Ok(())
674    }
675
676    fn open(&mut self, end: End, suffix: Suffix, memoize: Option<u64>, is_map: bool, in_key: bool) {
677        self.stack.push(Open {
678            end,
679            suffix,
680            memoize,
681            is_map,
682            expects_key: true,
683            in_key,
684        });
685    }
686
687    fn close(&mut self) -> Result<(), Error> {
688        let open = self.stack.pop().unwrap();
689        self.out.extend_from_slice(match open.end {
690            End::Appends => b"e",
691            End::SetItems => b"u",
692            End::AddItems => b"\x90",
693            End::Tuple => b"t",
694            End::FrozenSet => b"\x91",
695            End::List => b"l",
696            End::Empty => b"",
697        });
698        self.out.extend_from_slice(open.suffix.as_slice());
699        self.memoize(open.memoize);
700        self.value_done();
701        Ok(())
702    }
703
704    /// The module of the builtins for the protocol.
705    fn builtins(&self) -> &'static str {
706        match self.proto {
707            2 => "__builtin__",
708            _ => "builtins",
709        }
710    }
711
712    fn memoize(&mut self, id: Option<u64>) {
713        let Some(id) = id else {
714            return;
715        };
716        let idx = self.memo_len;
717        self.memo_len += 1;
718        self.memo.insert(id, idx);
719        if self.proto >= 4 {
720            self.out.push(0x94);
721        } else if idx < 256 {
722            self.out.push(b'q');
723            self.out.push(idx as u8);
724        } else {
725            self.out.push(b'r');
726            self.out.extend_from_slice(&idx.to_le_bytes());
727        }
728    }
729
730    fn get(&mut self, idx: u32) {
731        if idx < 256 {
732            self.out.push(b'h');
733            self.out.push(idx as u8);
734        } else {
735            self.out.push(b'j');
736            self.out.extend_from_slice(&idx.to_le_bytes());
737        }
738    }
739
740    fn global(&mut self, global: &Global) -> Result<(), Error> {
741        let (mut module, mut name) = (global.module(), global.name());
742        // like Python, the globals of Python 3 are written with the names
743        // of Python 2 before protocol 3
744        if self.proto < 3
745            && let Some((new_module, new_name)) = compat::reverse_fix_import(module, name)
746        {
747            module = new_module;
748            name = new_name.unwrap_or(name);
749        }
750        if self.proto >= 4 {
751            self.string(module)?;
752            self.string(name)?;
753            self.out.push(0x93);
754        } else {
755            let valid = |x: &str| !x.is_empty() && !x.contains('\n');
756            if !valid(module) || !valid(name) {
757                return Err(Error::new(
758                    ErrorKind::InvalidValue,
759                    "names of globals must not be empty or contain newlines before protocol 4",
760                ));
761            }
762            self.out.push(b'c');
763            self.out.extend_from_slice(module.as_bytes());
764            self.out.push(b'\n');
765            self.out.extend_from_slice(name.as_bytes());
766            self.out.push(b'\n');
767        }
768        Ok(())
769    }
770
771    fn atom(&mut self, atom: Atom<'_>, kind: Option<Kind>) -> Result<(), Error> {
772        match atom {
773            Atom::Null => self.out.push(b'N'),
774            Atom::Bool(value) => self.out.push(if value { 0x88 } else { 0x89 }),
775            Atom::U64(value) => self.int(value.into()),
776            Atom::I64(value) => self.int(value.into()),
777            Atom::F32(value) => self.float(value.into()),
778            Atom::F64(value) => self.float(value),
779            Atom::Str(value) | Atom::Lexical(value) => self.string(value.as_str())?,
780            Atom::Char(value) => self.string(value.encode_utf8(&mut [0; 4]))?,
781            Atom::Bytes(value) => match kind {
782                Some(Kind::ByteArray) => self.bytearray(value.data())?,
783                _ => self.bytes(value.data())?,
784            },
785            // values whose type was inferred from text are written as value
786            Atom::Implicit(value) => return self.atom(value.value().to_atom(), kind),
787            Atom::Ext(ref ext) => return self.ext(ext, kind),
788            _ => return Err(Error::new(ErrorKind::UnsupportedType, "unknown atom")),
789        }
790        Ok(())
791    }
792
793    #[cold]
794    fn ext(&mut self, ext: &ExtValue<'_>, kind: Option<Kind>) -> Result<(), Error> {
795        if let Some(reference) = ext.downcast_ref::<Reference>() {
796            return match self.memo.get(&reference.id()) {
797                Some(&idx) => {
798                    self.get(idx);
799                    Ok(())
800                }
801                None => Err(Error::new(
802                    ErrorKind::InvalidValue,
803                    "reference to a value that was not written or a tuple that contains itself",
804                )),
805            };
806        }
807        if let Some(global) = ext.downcast_ref::<Global>() {
808            return self.global(global);
809        }
810        if let Some(&value) = ext.downcast_ref::<u128>() {
811            self.long(&BigInt {
812                negative: false,
813                magnitude: value.to_be_bytes().to_vec(),
814            });
815            return Ok(());
816        }
817        if let Some(&value) = ext.downcast_ref::<i128>() {
818            self.long(&BigInt {
819                negative: value < 0,
820                magnitude: value.unsigned_abs().to_be_bytes().to_vec(),
821            });
822            return Ok(());
823        }
824        if let Some(value) = ext.downcast_ref::<BigInt>() {
825            self.long(value);
826            return Ok(());
827        }
828        // numbers from text formats are integers if their text is one
829        if let Some(value) = ext.downcast_value_ref::<Number>() {
830            return match value.as_str().parse::<i64>() {
831                Ok(value) => {
832                    self.int(value.into());
833                    Ok(())
834                }
835                Err(_) => match value.as_str().parse::<BigInt>() {
836                    Ok(value) => {
837                        self.long(&value);
838                        Ok(())
839                    }
840                    Err(_) => {
841                        self.float(value.value());
842                        Ok(())
843                    }
844                },
845            };
846        }
847        match ext.fallback() {
848            Atom::Ext(_) => Err(Error::new(
849                ErrorKind::UnsupportedType,
850                format!("pickle does not support {}", ext.name()),
851            )),
852            fallback => self.atom(fallback, kind),
853        }
854    }
855
856    fn int(&mut self, value: i128) {
857        if let Ok(value) = u8::try_from(value) {
858            self.out.push(b'K');
859            self.out.push(value);
860        } else if let Ok(value) = u16::try_from(value) {
861            self.out.push(b'M');
862            self.out.extend_from_slice(&value.to_le_bytes());
863        } else if let Ok(value) = i32::try_from(value) {
864            self.out.push(b'J');
865            self.out.extend_from_slice(&value.to_le_bytes());
866        } else {
867            self.long(&BigInt {
868                negative: value < 0,
869                magnitude: value.unsigned_abs().to_be_bytes().to_vec(),
870            });
871        }
872    }
873
874    /// Writes an integer with `LONG1` or `LONG4` (little-endian two's
875    /// complement).
876    fn long(&mut self, value: &BigInt) {
877        let magnitude = value.significant_magnitude();
878        // one more byte than the magnitude has room for the sign
879        let mut bytes: Vec<u8> = magnitude.iter().rev().copied().collect();
880        bytes.push(0);
881        if value.is_negative() {
882            for byte in bytes.iter_mut() {
883                *byte = !*byte;
884            }
885            for byte in bytes.iter_mut() {
886                let (sum, overflow) = byte.overflowing_add(1);
887                *byte = sum;
888                if !overflow {
889                    break;
890                }
891            }
892        }
893        // drop the bytes that only repeat the sign
894        while bytes.len() > 1 {
895            let last = bytes[bytes.len() - 1];
896            let before = bytes[bytes.len() - 2];
897            if (last == 0 && before & 0x80 == 0) || (last == 0xff && before & 0x80 != 0) {
898                bytes.pop();
899            } else {
900                break;
901            }
902        }
903        if value.is_zero() {
904            bytes.clear();
905        }
906        if bytes.len() < 256 {
907            self.out.push(0x8a);
908            self.out.push(bytes.len() as u8);
909        } else {
910            self.out.push(0x8b);
911            self.out
912                .extend_from_slice(&(bytes.len() as u32).to_le_bytes());
913        }
914        self.out.extend_from_slice(&bytes);
915    }
916
917    fn float(&mut self, value: f64) {
918        self.out.push(b'G');
919        self.out.extend_from_slice(&value.to_be_bytes());
920    }
921
922    fn string(&mut self, value: &str) -> Result<(), Error> {
923        let len = value.len();
924        if self.proto >= 4 && len < 256 {
925            self.out.push(0x8c);
926            self.out.push(len as u8);
927        } else if let Ok(len) = u32::try_from(len) {
928            self.out.push(b'X');
929            self.out.extend_from_slice(&len.to_le_bytes());
930        } else if self.proto >= 4 {
931            self.out.push(0x8d);
932            self.out.extend_from_slice(&(len as u64).to_le_bytes());
933        } else {
934            return Err(too_large());
935        }
936        self.out.extend_from_slice(value.as_bytes());
937        Ok(())
938    }
939
940    fn bytes(&mut self, value: &[u8]) -> Result<(), Error> {
941        if self.proto < 3 {
942            // Python 3 writes `_codecs.encode(text, "latin1")`
943            if value.is_empty() {
944                self.out.extend_from_slice(b"c__builtin__\nbytes\n)R");
945            } else {
946                self.out.extend_from_slice(b"c_codecs\nencode\n");
947                self.string(&latin1(value))?;
948                self.string("latin1")?;
949                self.out.extend_from_slice(b"\x86R");
950            }
951            return Ok(());
952        }
953        let len = value.len();
954        if len < 256 {
955            self.out.push(b'C');
956            self.out.push(len as u8);
957        } else if let Ok(len) = u32::try_from(len) {
958            self.out.push(b'B');
959            self.out.extend_from_slice(&len.to_le_bytes());
960        } else if self.proto >= 4 {
961            self.out.push(0x8e);
962            self.out.extend_from_slice(&(len as u64).to_le_bytes());
963        } else {
964            return Err(too_large());
965        }
966        self.out.extend_from_slice(value);
967        Ok(())
968    }
969
970    fn bytearray(&mut self, value: &[u8]) -> Result<(), Error> {
971        match self.proto {
972            5.. => {
973                self.out.push(0x96);
974                self.out
975                    .extend_from_slice(&(value.len() as u64).to_le_bytes());
976                self.out.extend_from_slice(value);
977            }
978            3 | 4 => {
979                self.global(&Global::new("builtins", "bytearray"))?;
980                self.bytes(value)?;
981                self.out.extend_from_slice(b"\x85R");
982            }
983            _ => {
984                self.out.extend_from_slice(b"c__builtin__\nbytearray\n");
985                self.string(&latin1(value))?;
986                self.string("latin-1")?;
987                self.out.extend_from_slice(b"\x86R");
988            }
989        }
990        Ok(())
991    }
992}
993
994/// Decodes bytes as latin-1.
995fn latin1(bytes: &[u8]) -> String {
996    bytes.iter().map(|&c| char::from(c)).collect()
997}
998
999#[cold]
1000fn too_large() -> Error {
1001    Error::new(
1002        ErrorKind::OutOfRange,
1003        "value too large for the pickle protocol",
1004    )
1005}
deser_pickle/ser.rs

deser_pickle/
ser.rs
