From 69781f0adf0053ae1604d95e2c9525f4d92df3e3 Mon Sep 17 00:00:00 2001
From: blead <blead010@gmail.com>
Date: Sat, 17 Aug 2024 23:14:59 +0700
Subject: [PATCH] Implement amf3 packing, fix amf3 extraction bug with empty
 arrays becoming null, rename flate functions

---
 pkg/encoding/amf3.go  | 58 ++++++++++++++++++++++++++++++++++++++++---
 pkg/encoding/flate.go | 25 ++++++++++++++++++-
 pkg/wf/packer.go      |  8 ++++--
 pkg/wf/parser.go      |  2 +-
 4 files changed, 85 insertions(+), 8 deletions(-)

diff --git a/pkg/encoding/amf3.go b/pkg/encoding/amf3.go
index f08d343..0d8c942 100644
--- a/pkg/encoding/amf3.go
+++ b/pkg/encoding/amf3.go
@@ -3,6 +3,7 @@ package encoding
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
 	"strings"
 
 	amf "github.com/remyoudompheng/goamf"
@@ -10,18 +11,18 @@ import (
 
 // Amf3ToJSON converts Adobe AMF3 to JSON.
 func Amf3ToJSON(raw []byte, indent int) ([]byte, error) {
-	data, err := deflate(raw)
+	data, err := inflate(raw)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("flate decompress error, %w", err)
 	}
 
 	d := amf.NewDecoder()
 	amf3, err := d.DecodeAmf3(bytes.NewReader(data))
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("amf3 decode error, %w", err)
 	}
 
-	js, err := jsonMarshalNoEscape(amf3)
+	js, err := jsonMarshalNoEscape(initAmfTypes(amf3))
 	if err != nil {
 		return nil, err
 	}
@@ -31,3 +32,52 @@ func Amf3ToJSON(raw []byte, indent int) ([]byte, error) {
 
 	return output.Bytes(), err
 }
+
+// JSONToAmf3 converts JSON to Adobe AMF3
+func JSONToAmf3(js []byte) ([]byte, error) {
+	var data any
+	err := json.Unmarshal(js, &data)
+	if err != nil {
+		return nil, err
+	}
+
+	var amf3 bytes.Buffer
+	e := new(amf.Encoder)
+	_, err = e.EncodeAmf3(&amf3, initAmfTypes(data))
+	if err != nil {
+		return nil, err
+	}
+
+	return deflate(amf3.Bytes())
+}
+
+// 1. empty arrays/maps returned from goamf.Decoder#DecodeAmf3 are nil pointers so they need to be initialized before marshalling to JSON
+// 2. primitive map[string]any needs to be converted to goamf.Object first for goamf.Encoder#EncodeAmf3 to work properly
+func initAmfTypes(p any) any {
+	switch v := p.(type) {
+	case map[string]any, amf.Object:
+		val, ok := v.(map[string]any)
+		if !ok {
+			val = map[string]any(v.(amf.Object))
+		}
+
+		obj := make(amf.Object)
+		for key, value := range val {
+			obj[key] = initAmfTypes(value)
+		}
+		return obj
+	case []any, amf.Array:
+		val, ok := v.([]any)
+		if !ok {
+			val = []any(v.(amf.Array))
+		}
+
+		arr := make(amf.Array, len(val))
+		for i, value := range val {
+			arr[i] = initAmfTypes(value)
+		}
+		return arr
+	default:
+		return v
+	}
+}
diff --git a/pkg/encoding/flate.go b/pkg/encoding/flate.go
index 43d4f2e..d0eb58b 100644
--- a/pkg/encoding/flate.go
+++ b/pkg/encoding/flate.go
@@ -6,7 +6,30 @@ import (
 	"io"
 )
 
-func deflate(compressed []byte) ([]byte, error) {
+// deflate compresses raw bytes
+func deflate(raw []byte) ([]byte, error) {
+	var output bytes.Buffer
+	zw, err := flate.NewWriter(&output, flate.DefaultCompression)
+	if err != nil {
+		return nil, err
+	}
+	defer zw.Close()
+
+	_, err = zw.Write(raw)
+	if err != nil {
+		return nil, err
+	}
+
+	err = zw.Close()
+	if err != nil {
+		return nil, err
+	}
+
+	return output.Bytes(), nil
+}
+
+// inflate decompresses deflate-encoded bytes
+func inflate(compressed []byte) ([]byte, error) {
 	zr := flate.NewReader(bytes.NewReader(compressed))
 	defer zr.Close()
 
diff --git a/pkg/wf/packer.go b/pkg/wf/packer.go
index 5b39900..5ebb5ee 100644
--- a/pkg/wf/packer.go
+++ b/pkg/wf/packer.go
@@ -68,8 +68,12 @@ func NewPacker(config *PackerConfig) (*Packer, error) {
 	}
 
 	parsers := []parser{
-		// &amf3Parser{ext: ".action.dsl"},
-		// &esdlParser{&amf3Parser{ext: ".esdl"}},
+		&amf3Parser{ext: ".action.dsl"},
+		&amf3Parser{ext: ".atlas"},
+		&amf3Parser{ext: ".frame"},
+		&amf3Parser{ext: ".parts"},
+		&amf3Parser{ext: ".timeline"},
+		&esdlParser{&amf3Parser{ext: ".esdl"}},
 		&pngParser{},
 		&orderedmapParser{}, // needs to be last because of ambiguous file extension
 	}
diff --git a/pkg/wf/parser.go b/pkg/wf/parser.go
index 10d36e4..ac1da08 100644
--- a/pkg/wf/parser.go
+++ b/pkg/wf/parser.go
@@ -78,7 +78,7 @@ func (parser *amf3Parser) matchDest(dest string, config *PackerConfig) (string,
 }
 
 func (*amf3Parser) unparse(raw []byte, config *PackerConfig) ([]byte, error) {
-	return nil, fmt.Errorf("unparse amf3Parser: not implemented")
+	return encoding.JSONToAmf3(raw)
 }
 
 type esdlParser struct {