1use crate::PADDING;
4use crate::amf3::custom_encoder::CustomEncoder;
5use crate::amf3::element_cache::ElementCache;
6use crate::amf3::length::Length;
7use crate::amf3::type_marker::TypeMarker;
8use crate::types::{
9 Attribute, ClassDefinition, DictionaryEntry, DictionaryObjectValue, Element, ObjectId, Value,
10 VectorObjectValue, VectorPrimitiveValue,
11};
12use crate::write::WriteExt;
13use std::cell::RefCell;
14use std::collections::{BTreeMap, HashMap};
15use std::io::Result;
16use std::io::Write;
17
18#[derive(Default)]
20pub struct AMF3Encoder {
21 string_reference_table: ElementCache<Vec<u8>>,
23
24 trait_reference_table: RefCell<Vec<ClassDefinition>>,
26
27 object_reference_table: ElementCache<Value>,
29
30 pub external_encoders: HashMap<String, Box<dyn CustomEncoder>>,
32
33 object_id_to_reference: RefCell<BTreeMap<ObjectId, (TypeMarker, usize)>>,
34}
35
36impl AMF3Encoder {
37 #[allow(clippy::unusual_byte_groupings)]
38 pub(crate) fn write_int<'a, 'b: 'a, W: Write + 'a>(
39 &self,
40 writer: &mut W,
41 i: i32,
42 ) -> Result<()> {
43 let n = if i < 0 {
44 i + 0b001_0000000_0000000_0000000_00000000
45 } else {
46 i
47 };
48
49 if n > 0x1fffff {
50 writer.write_u8(((n >> (7 * 3 + 1)) | 0b10000000) as u8)?;
51 writer.write_u8(((n >> (7 * 2 + 1)) | 0b10000000) as u8)?;
52 writer.write_u8(((n >> (7 + 1)) | 0b10000000) as u8)?;
53 writer.write_u8((n & 0b11111111) as u8)?;
54 } else if n > 0x3fff {
55 writer.write_u8(((n >> (7 * 2)) | 0b10000000) as u8)?;
56 writer.write_u8(((n >> 7) | 0b10000000) as u8)?;
57 writer.write_u8((n & 0b01111111) as u8)?;
58 } else if n > 0x7f {
59 writer.write_u8(((n >> 7) | 0b10000000) as u8)?;
60 writer.write_u8((n & 0b01111111) as u8)?;
61 } else {
62 writer.write_u8((n & 0b01111111) as u8)?;
63 }
64
65 Ok(())
66 }
67
68 fn write_byte_string<'a, 'b: 'a, W: Write + 'a>(
69 &'a self,
70 writer: &mut W,
71 s: &'b [u8],
72 ) -> Result<()> {
73 let len = if !s.is_empty() {
74 self.string_reference_table
75 .to_length(s.to_vec(), s.len() as u32)
76 } else {
77 Length::Size(0)
78 };
79
80 let only_length = len.is_reference() && !s.is_empty();
81 let s_vec = s.to_vec();
82
83 if !s_vec.is_empty() {
84 self.string_reference_table.store(s_vec.clone());
85 }
86
87 len.write(writer, self)?;
88 if !only_length {
89 writer.write_all(s)?;
90 }
91
92 Ok(())
93 }
94
95 fn write_string<'a, 'b: 'a, W: Write + 'a>(&'a self, writer: &mut W, s: &'b str) -> Result<()> {
96 self.write_byte_string(writer, s.as_bytes())
97 }
98
99 fn write_type_marker<'a, 'b: 'a, W: Write + 'a>(
100 &self,
101 writer: &mut W,
102 s: TypeMarker,
103 ) -> Result<()> {
104 writer.write_u8(s as u8)
105 }
106
107 fn write_number_element<'a, 'b: 'a, W: Write + 'a>(
108 &self,
109 writer: &mut W,
110 i: f64,
111 ) -> Result<()> {
112 self.write_type_marker(writer, TypeMarker::Number)?;
113 writer.write_f64(i)?;
114 Ok(())
115 }
116
117 fn write_boolean_element<'a, 'b: 'a, W: Write + 'a>(
118 &self,
119 writer: &mut W,
120 b: bool,
121 ) -> Result<()> {
122 if b {
123 self.write_type_marker(writer, TypeMarker::True)
124 } else {
125 self.write_type_marker(writer, TypeMarker::False)
126 }
127 }
128
129 fn write_string_element<'a, 'b: 'a, W: Write + 'a>(
130 &'a self,
131 writer: &mut W,
132 s: &'b str,
133 ) -> Result<()> {
134 self.write_type_marker(writer, TypeMarker::String)?;
135 self.write_byte_string(writer, s.as_bytes())?;
136 Ok(())
137 }
138
139 fn write_null_element<'a, 'b: 'a, W: Write + 'a>(&self, writer: &mut W) -> Result<()> {
140 self.write_type_marker(writer, TypeMarker::Null)
141 }
142
143 fn write_undefined_element<'a, 'b: 'a, W: Write + 'a>(&self, writer: &mut W) -> Result<()> {
144 self.write_type_marker(writer, TypeMarker::Undefined)
145 }
146
147 fn write_int_vector<'a, 'b: 'a, W: Write + 'a>(
148 &'a self,
149 writer: &mut W,
150 items: &'b [i32],
151 fixed_length: bool,
152 ) -> Result<()> {
153 let len = self.object_reference_table.to_length(
154 Value::VectorInt(VectorPrimitiveValue {
155 values: items.to_vec(),
156 fixed_length,
157 }),
158 items.len() as u32,
159 );
160
161 self.write_type_marker(writer, TypeMarker::VectorInt)?;
162 if len.is_reference() {
163 len.write(writer, self)?;
164 } else {
165 Length::Size(items.len() as u32).write(writer, self)?;
166 writer.write_u8(fixed_length as u8)?;
167 for item in items {
168 writer.write_i32(*item)?;
169 }
170 }
171 Ok(())
172 }
173
174 fn write_uint_vector<'a, 'b: 'a, W: Write + 'a>(
175 &'a self,
176 writer: &mut W,
177 items: &'b [u32],
178 fixed_length: bool,
179 ) -> Result<()> {
180 let len = self.object_reference_table.to_length(
181 Value::VectorUInt(VectorPrimitiveValue {
182 values: items.to_vec(),
183 fixed_length,
184 }),
185 items.len() as u32,
186 );
187
188 self.write_type_marker(writer, TypeMarker::VectorUInt)?;
189 if len.is_reference() {
190 len.write(writer, self)?;
191 } else {
192 Length::Size(items.len() as u32).write(writer, self)?;
193 writer.write_u8(fixed_length as u8)?;
194 for item in items {
195 writer.write_u32(*item)?;
196 }
197 }
198 Ok(())
199 }
200
201 fn write_number_vector<'a, 'b: 'a, W: Write + 'a>(
202 &'a self,
203 writer: &mut W,
204 items: &'b [f64],
205 fixed_length: bool,
206 ) -> Result<()> {
207 let len = self.object_reference_table.to_length(
208 Value::VectorDouble(VectorPrimitiveValue {
209 values: items.to_vec(),
210 fixed_length,
211 }),
212 items.len() as u32,
213 );
214
215 self.write_type_marker(writer, TypeMarker::VectorDouble)?;
216 if len.is_reference() {
217 len.write(writer, self)?;
218 } else {
219 Length::Size(items.len() as u32).write(writer, self)?;
220 writer.write_u8(fixed_length as u8)?;
221 for item in items {
222 writer.write_f64(*item)?;
223 }
224 }
225 Ok(())
226 }
227
228 fn write_date_element<'a, 'b: 'a, W: Write + 'a>(
229 &'a self,
230 writer: &mut W,
231 time: f64,
232 ) -> Result<()> {
233 let len = self.object_reference_table.to_length(
234 Value::Date {
235 time,
236 timezone_or_utc: None,
237 },
238 0,
239 );
240
241 self.write_type_marker(writer, TypeMarker::Date)?;
242 len.write(writer, self)?;
243 if len.is_size() {
244 writer.write_f64(time)?;
245 }
246 Ok(())
247 }
248
249 fn write_integer_element<'a, 'b: 'a, W: Write + 'a>(
250 &self,
251 writer: &mut W,
252 i: i32,
253 ) -> Result<()> {
254 self.write_type_marker(writer, TypeMarker::Integer)?;
255 self.write_int(writer, i)?;
256 Ok(())
257 }
258
259 fn write_byte_array_element<'a, 'b: 'a, W: Write + 'a>(
260 &'a self,
261 writer: &mut W,
262 bytes: &'b [u8],
263 ) -> Result<()> {
264 let len = self
265 .object_reference_table
266 .to_length(Value::ByteArray(bytes.to_vec()), bytes.len() as u32);
267
268 self.write_type_marker(writer, TypeMarker::ByteArray)?;
269 len.write(writer, self)?;
270 if len.is_size() {
271 writer.write_all(bytes)?;
272 }
273 Ok(())
274 }
275
276 fn write_xml_element<'a, 'b: 'a, W: Write + 'a>(
277 &self,
278 writer: &mut W,
279 bytes: &'b str,
280 string: bool,
281 ) -> Result<()> {
282 let len = Length::Size(bytes.len() as u32);
283
284 if string {
285 self.write_type_marker(writer, TypeMarker::XmlString)?;
286 } else {
287 self.write_type_marker(writer, TypeMarker::Xml)?;
288 }
289
290 len.write(writer, self)?;
291 if len.is_size() {
292 writer.write_all(bytes.as_bytes())?;
293 }
294 Ok(())
295 }
296
297 fn write_class_definition<'a, 'b: 'a, W: Write + 'a>(
298 &'a self,
299 writer: &mut W,
300 class_def: &'b ClassDefinition,
301 ) -> Result<()> {
302 self.write_byte_string(writer, class_def.name.as_bytes())?;
303 for p in &class_def.static_properties {
304 self.write_string(writer, p)?;
305 }
306 Ok(())
307 }
308
309 fn write_trait_reference<'a, 'b: 'a, W: Write + 'a>(
311 &'a self,
312 writer: &mut W,
313 index: u32,
314 children: &'b [Element],
315 custom_props: Option<&'b [Element]>,
316 def: &'b ClassDefinition,
317 ) -> Result<()> {
318 #[allow(clippy::identity_op)]
319 let size = (((index << 1) | 0u32) << 1) | 1u32;
320
321 self.write_int(writer, size as i32)?;
322 if def.attributes.contains(Attribute::External) {
323 if let Some(encoder) = self.external_encoders.get(&def.name) {
324 writer.write_all(&encoder.encode(
325 custom_props.expect("Custom encoder, missing custom props"),
326 &Some(def.clone()),
327 self,
328 ))?;
329 } else {
330 unimplemented!();
331 }
332 }
333
334 if !def.attributes.contains(Attribute::External) {
335 if def.attributes.is_empty() {
336 for c in children {
337 if def.static_properties.contains(&c.name) {
338 self.write_value_element(writer, &c.value)?;
339 }
340 }
341 }
342
343 if def.attributes.contains(Attribute::Dynamic) {
344 for c in children {
345 if def.static_properties.contains(&c.name) {
346 self.write_value_element(writer, &c.value)?;
347 }
348 }
349
350 for c in children {
351 if !def.static_properties.contains(&c.name) {
352 self.write_byte_string(writer, c.name.as_bytes())?;
353 self.write_value_element(writer, &c.value)?;
354 }
355 }
356
357 self.write_byte_string(writer, &[])?;
358 }
359 }
360 Ok(())
361 }
362
363 fn write_object_reference<'a, 'b: 'a, W: Write + 'a>(
364 &'a self,
365 writer: &mut W,
366 index: u32,
367 ) -> Result<()> {
368 #[allow(clippy::identity_op)]
369 let size = (index << 1) | 0u32;
370 self.write_int(writer, size as i32)
371 }
372
373 fn write_object_full<'a, 'b: 'a, W: Write + 'a>(
374 &'a self,
375 writer: &mut W,
376 custom_props: Option<&'b [Element]>,
377 children: &'b [Element],
378 def: &'b ClassDefinition,
379 ) -> Result<()> {
380 let is_external = def.attributes.contains(Attribute::External);
381 let is_dynamic = def.attributes.contains(Attribute::Dynamic);
382
383 let mut encoding = 0b00;
384 if is_external {
385 encoding |= 0b01;
386 }
387 if is_dynamic {
388 encoding |= 0b10;
389 }
390
391 let size = ((((((def.static_properties.len() as u32) << 2) | (encoding & 0xff) as u32)
393 << 1)
394 | 1u32)
395 << 1)
396 | 1u32;
397
398 self.trait_reference_table.borrow_mut().push(def.clone());
399 self.write_int(writer, size as i32)?;
400 self.write_class_definition(writer, def)?;
401 if def.attributes.contains(Attribute::External) {
402 if let Some(encoder) = self.external_encoders.get(&def.name) {
403 writer.write_all(&encoder.encode(
404 custom_props.expect("Custom encode but missing custom props"),
405 &Some(def.clone()),
406 self,
407 ))?;
408 } else {
409 unimplemented!();
410 }
411 }
412 if !def.attributes.contains(Attribute::External) {
413 if def.attributes.is_empty() {
414 for c in children {
415 if def.static_properties.contains(&c.name) {
416 self.write_value_element(writer, &c.value)?;
417 }
418 }
419 }
420
421 if def.attributes.contains(Attribute::Dynamic) {
422 for c in children {
423 if def.static_properties.contains(&c.name) {
424 self.write_value_element(writer, &c.value)?;
425 }
426 }
427 for c in children {
428 if !def.static_properties.contains(&c.name) {
429 self.write_byte_string(writer, c.name.as_bytes())?;
430 self.write_value_element(writer, &c.value)?;
431 }
432 }
433 self.write_byte_string(writer, &[])?;
434 }
435 }
436 Ok(())
437 }
438
439 fn write_object_element<'a, 'b: 'a, W: Write + 'a>(
440 &'a self,
441 writer: &mut W,
442 id: ObjectId,
443 children: &'b [Element],
444 custom_props: Option<&'b [Element]>,
445 class_def: &'b Option<ClassDefinition>,
446 ) -> Result<()> {
447 let had_object = Length::Size(0);
448 use crate::types::ObjectValue;
449
450 let obj = Value::Object {
451 id,
452 data: ObjectValue {
453 elements: children.to_vec(),
454 class_definition: class_def.clone(),
455 },
456 };
457 self.object_reference_table.store(obj.clone());
458 if let Length::Reference(r) = self.object_reference_table.to_length(obj, 0) {
459 self.object_id_to_reference
460 .borrow_mut()
461 .insert(id, (TypeMarker::Object, r));
462 }
463
464 let def = class_def.clone().unwrap_or_default();
465 let def2 = def.clone();
466
467 let has_trait = self
468 .trait_reference_table
469 .borrow()
470 .iter()
471 .position(|cd| *cd == def);
472
473 self.write_type_marker(writer, TypeMarker::Object)?;
474 if had_object.is_reference() {
475 self.write_object_reference(
476 writer,
477 had_object
478 .as_position()
479 .expect("Failed to convert to position") as u32,
480 )?;
481 }
482 if !had_object.is_reference() {
483 if let Some(has_trait) = has_trait {
484 self.write_trait_reference(
485 writer,
486 has_trait as u32,
487 children,
488 custom_props,
489 &def2,
490 )?;
491 }
492 if has_trait.is_none() {
493 self.write_object_full(writer, custom_props, children, &def)?;
494 }
495 }
496
497 Ok(())
498 }
499
500 fn write_strict_array_element<'a, 'b: 'a, W: Write + 'a>(
501 &'a self,
502 writer: &mut W,
503 children: &'b [Value],
504 ) -> Result<()> {
505 let len = Length::Size(children.len() as u32);
507
508 if children.is_empty() {
510 self.write_type_marker(writer, TypeMarker::Array)?;
511 Length::Size(0).write(writer, self)?;
512 self.write_byte_string(writer, &[])?; } else {
514 self.write_type_marker(writer, TypeMarker::Array)?;
515 len.write(writer, self)?;
516 if len.is_size() {
517 self.write_byte_string(writer, &[])?; for v in children {
519 self.write_value_element(writer, v)?;
520 }
521 }
522 }
523
524 Ok(())
525 }
526
527 fn write_ecma_array_element<'a, 'b: 'a, W: Write + 'a>(
528 &'a self,
529 writer: &mut W,
530 dense: &'b [Value],
531 assoc: &'b [Element],
532 ) -> Result<()> {
533 let len = Length::Size(dense.len() as u32);
534
535 self.write_type_marker(writer, TypeMarker::Array)?;
538 len.write(writer, self)?;
539 if len.is_size() {
540 for out in assoc {
541 self.write_element(writer, out)?;
542 }
543 self.write_byte_string(writer, &[])?;
544 for out in dense {
545 self.write_value_element(writer, out)?;
546 }
547 }
548 Ok(())
549 }
550
551 fn write_object_vector_element<'a, 'b: 'a, W: Write + 'a>(
552 &'a self,
553 writer: &mut W,
554 id: ObjectId,
555 items: &'b [Value],
556 type_name: &'b str,
557 fixed_length: bool,
558 ) -> Result<()> {
559 let vo = Value::VectorObject {
560 id,
561 data: VectorObjectValue {
562 values: items.to_vec(),
563 object_type_name: type_name.to_string(),
564 fixed_length,
565 },
566 };
567
568 let len = self
569 .object_reference_table
570 .to_length(vo.clone(), items.len() as u32);
571 self.object_reference_table.store(vo.clone());
572 if let Length::Reference(r) = self.object_reference_table.to_length(vo, 0) {
573 self.object_id_to_reference
574 .borrow_mut()
575 .insert(id, (TypeMarker::VectorObject, r));
576 }
577
578 self.write_type_marker(writer, TypeMarker::VectorObject)?;
579 len.write(writer, self)?;
580 if len.is_size() {
581 writer.write_u8(fixed_length as u8)?;
582 self.write_string(writer, type_name)?;
583 for i in items {
584 self.write_value_element(writer, i)?;
585 }
586 }
587 Ok(())
588 }
589
590 fn write_dictionary_element<'a, 'b: 'a, W: Write + 'a>(
591 &'a self,
592 writer: &mut W,
593 id: ObjectId,
594 items: &'b [DictionaryEntry],
595 weak_keys: bool,
596 ) -> Result<()> {
597 let dict = Value::Dictionary {
598 id,
599 data: DictionaryObjectValue {
600 elements: items.to_vec(),
601 weak_keys,
602 },
603 };
604
605 let len = self
606 .object_reference_table
607 .to_length(dict.clone(), items.len() as u32);
608 self.object_reference_table.store(dict.clone());
609 if let Length::Reference(r) = self.object_reference_table.to_length(dict, 0) {
610 self.object_id_to_reference
611 .borrow_mut()
612 .insert(id, (TypeMarker::Dictionary, r));
613 }
614
615 self.write_type_marker(writer, TypeMarker::Dictionary)?;
616 len.write(writer, self)?;
617 if len.is_size() {
618 writer.write_u8(weak_keys as u8)?;
619 for i in items {
620 self.write_value_element(writer, &i.key)?;
621 self.write_value_element(writer, &i.value)?;
622 }
623 }
624 Ok(())
625 }
626
627 pub(crate) fn write_value_element<'a, 'b: 'a, W: Write + 'a>(
628 &'b self,
629 writer: &mut W,
630 s: &'b Value,
631 ) -> Result<()> {
632 self.write_value(writer, s)
633 }
634
635 fn write_value<'a, 'b: 'a, W: Write + 'a>(
636 &'b self,
637 writer: &mut W,
638 s: &'b Value,
639 ) -> Result<()> {
640 match s {
641 Value::Number(x) => self.write_number_element(writer, *x),
642 Value::Bool(b) => self.write_boolean_element(writer, *b),
643 Value::String(s) => self.write_string_element(writer, s),
644 Value::Object { id, data } => {
645 self.write_object_element(writer, *id, &data.elements, None, &data.class_definition)
646 }
647 Value::Null => self.write_null_element(writer),
648 Value::Undefined => self.write_undefined_element(writer),
649 Value::ECMAArray { id, data } => {
650 self.object_reference_table.store(s.clone());
651 if let Length::Reference(r) = self.object_reference_table.to_length(s.clone(), 0) {
652 self.object_id_to_reference
653 .borrow_mut()
654 .insert(*id, (TypeMarker::Array, r));
655 }
656
657 self.write_ecma_array_element(writer, &data.dense, &data.elements)
658 }
659 Value::StrictArray { id, values } => {
660 self.object_reference_table.store(s.clone());
661 if let Length::Reference(r) = self.object_reference_table.to_length(s.clone(), 0) {
662 self.object_id_to_reference
663 .borrow_mut()
664 .insert(*id, (TypeMarker::Array, r));
665 }
666
667 self.write_strict_array_element(writer, values)
668 }
669 Value::Date {
670 time,
671 timezone_or_utc: _,
672 } => self.write_date_element(writer, *time),
673 Value::XML { value, is_string } => self.write_xml_element(writer, value, *is_string),
674 Value::Integer(i) => self.write_integer_element(writer, *i),
675 Value::ByteArray(bytes) => self.write_byte_array_element(writer, bytes),
676 Value::VectorInt(v) => self.write_int_vector(writer, &v.values, v.fixed_length),
677 Value::VectorUInt(v) => self.write_uint_vector(writer, &v.values, v.fixed_length),
678 Value::VectorDouble(v) => self.write_number_vector(writer, &v.values, v.fixed_length),
679 Value::VectorObject { id, data } => self.write_object_vector_element(
680 writer,
681 *id,
682 &data.values,
683 &data.object_type_name,
684 data.fixed_length,
685 ),
686 Value::Dictionary { id, data } => {
687 self.write_dictionary_element(writer, *id, &data.elements, data.weak_keys)
688 }
689
690 Value::Custom(c) => self.write_object_element(
691 writer,
692 ObjectId::INVALID,
693 &c.dynamic_elements,
694 Some(&c.elements),
695 &Some(c.class_definition.clone()),
696 ),
697 Value::AMF3(e) => self.write_value_element(writer, e),
698 Value::Unsupported => self.write_undefined_element(writer),
699 Value::Reference(_) => unimplemented!(),
700 Value::Amf3ObjectReference(id) => {
701 let (ty, r) = *self
702 .object_id_to_reference
703 .borrow()
704 .get(id)
705 .expect("Invalid reference");
706 self.write_type_marker(writer, ty)?;
707 self.write_object_reference(writer, r as u32)
708 }
709 }
710 }
711
712 fn write_element<'a, 'b: 'a, W: Write + 'a>(
713 &'b self,
714 writer: &mut W,
715 element: &'b Element,
716 ) -> Result<()> {
717 self.write_string(writer, &element.name)?;
718 self.write_value_element(writer, &element.value)?;
719 Ok(())
720 }
721
722 fn write_element_and_padding<'a, 'b: 'a, W: Write + 'a>(
723 &'b self,
724 writer: &mut W,
725 element: &'b Element,
726 ) -> Result<()> {
727 self.write_element(writer, element)?;
728 writer.write_all(&PADDING)?;
729 Ok(())
730 }
731
732 pub(crate) fn write_body<'a, 'b: 'a, W: Write + 'a>(
733 &'b self,
734 writer: &mut W,
735 elements: &'b [Element],
736 ) -> Result<()> {
737 for e in elements {
738 self.write_element_and_padding(writer, e)?;
739 }
740 Ok(())
741 }
742}
743
744#[cfg(test)]
745mod write_number_tests {
746 use crate::amf3::write::AMF3Encoder;
747
748 #[test]
749 fn test_write_1byte_number() {
750 let e = AMF3Encoder::default();
751 let mut v = Vec::new();
752 e.write_int(&mut v, 0b00101011).expect("Test fail");
753 assert_eq!(v, &[0b00101011]);
754 }
755
756 #[test]
757 fn test_write_4byte_number() {
758 let e = AMF3Encoder::default();
759 let mut v = Vec::new();
760 e.write_int(&mut v, 2097280).expect("Test fail");
761 assert_eq!(v, &[0b10000000, 0b11000000, 0b10000000, 0b10000000]);
762 }
763
764 #[test]
765 fn write_neg_number() {
766 let e = AMF3Encoder::default();
767 let mut v = Vec::new();
768 e.write_int(&mut v, -268435455).expect("Test fail");
769 assert_eq!(v, &[192, 128, 128, 1]);
770 }
771}