Fix marshaling master element from chan (#111)

at-wat · Dec 20, 2019 · cb2a64f · cb2a64f
1 parent 5acd530
commit cb2a64f
Show file tree

Hide file tree

Showing 2 changed files with 79 additions and 24 deletions.
diff --git a/marshal.go b/marshal.go
@@ -225,7 +225,14 @@ func marshalImpl(vo reflect.Value, w io.Writer, pos uint64, parent *Element, opt
 					if !ok {
 						break
 					}
-					pos, err = writeOne(val)
+					lst, ok := pealElem(val, e.t == DataTypeBinary, tag.omitEmpty)
+					if !ok {
+						return pos, ErrIncompatibleType
+					}
+					if len(lst) != 1 {
+						return pos, ErrIncompatibleType
+					}
+					pos, err = writeOne(lst[0])
 				}
 			default:
 				pos, err = writeOne(vn)

diff --git a/marshal_test.go b/marshal_test.go
@@ -406,37 +406,85 @@ func TestMarshal_InvalidTag(t *testing.T) {
 }
 
 func TestMarshal_Chan(t *testing.T) {
-	type Cluster struct {
-		Timecode uint64 `ebml:"Timecode"`
-	}
-	type TestOmitempty struct {
-		Segment struct {
-			Cluster chan Cluster `ebml:"Cluster,size=unknown"`
-		} `ebml:"Segment,size=unknown"`
-	}
-
-	ch := make(chan Cluster, 100)
-	input := &TestOmitempty{}
-	input.Segment.Cluster = ch
-	ch <- Cluster{Timecode: 0x01}
-	ch <- Cluster{Timecode: 0x02}
-	close(ch)
-
 	expected := []byte{
 		0x18, 0x53, 0x80, 0x67, 0x01, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
 		0x1F, 0x43, 0xB6, 0x75, 0x01, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
 		0xE7, 0x81, 0x01,
 		0x1F, 0x43, 0xB6, 0x75, 0x01, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
 		0xE7, 0x81, 0x02,
 	}
-
-	var b bytes.Buffer
-	if err := Marshal(input, &b); err != nil {
-		t.Fatalf("Unexpected error: %+v", err)
-	}
-	if !bytes.Equal(expected, b.Bytes()) {
-		t.Errorf("Marshaled binary doesn't match:\n expected: %v,\n      got: %v", expected, b.Bytes())
+	type Cluster struct {
+		Timecode uint64 `ebml:"Timecode"`
 	}
+
+	t.Run("ChanStruct", func(t *testing.T) {
+		ch := make(chan Cluster, 100)
+		input := &struct {
+			Segment struct {
+				Cluster chan Cluster `ebml:"Cluster,size=unknown"`
+			} `ebml:"Segment,size=unknown"`
+		}{}
+		input.Segment.Cluster = ch
+		ch <- Cluster{Timecode: 0x01}
+		ch <- Cluster{Timecode: 0x02}
+		close(ch)
+
+		var b bytes.Buffer
+		if err := Marshal(input, &b); err != nil {
+			t.Fatalf("Unexpected error: %+v", err)
+		}
+		if !bytes.Equal(expected, b.Bytes()) {
+			t.Errorf("Marshaled binary doesn't match:\n expected: %v,\n      got: %v", expected, b.Bytes())
+		}
+	})
+	t.Run("ChanStructPtr", func(t *testing.T) {
+		input := &struct {
+			Segment struct {
+				Cluster chan *Cluster `ebml:"Cluster,size=unknown"`
+			} `ebml:"Segment,size=unknown"`
+		}{}
+
+		t.Run("Valid", func(t *testing.T) {
+			ch := make(chan *Cluster, 100)
+			input.Segment.Cluster = ch
+			ch <- &Cluster{Timecode: 0x01}
+			ch <- &Cluster{Timecode: 0x02}
+			close(ch)
+
+			var b bytes.Buffer
+			if err := Marshal(input, &b); err != nil {
+				t.Fatalf("Unexpected error: %+v", err)
+			}
+			if !bytes.Equal(expected, b.Bytes()) {
+				t.Errorf("Marshaled binary doesn't match:\n expected: %v,\n      got: %v", expected, b.Bytes())
+			}
+		})
+		t.Run("Nil", func(t *testing.T) {
+			ch := make(chan *Cluster, 100)
+			input.Segment.Cluster = ch
+			ch <- nil
+			close(ch)
+
+			if err := Marshal(input, &bytes.Buffer{}); err != ErrIncompatibleType {
+				t.Fatalf("Expected %v, got %v", ErrIncompatibleType, err)
+			}
+		})
+	})
+	t.Run("ChanStructSlice", func(t *testing.T) {
+		input := &struct {
+			Segment struct {
+				Cluster chan []Cluster `ebml:"Cluster,size=unknown"`
+			} `ebml:"Segment,size=unknown"`
+		}{}
+		ch := make(chan []Cluster, 100)
+		input.Segment.Cluster = ch
+		ch <- make([]Cluster, 2)
+		close(ch)
+
+		if err := Marshal(input, &bytes.Buffer{}); err != ErrIncompatibleType {
+			t.Fatalf("Expected %v, got %v", ErrIncompatibleType, err)
+		}
+	})
 }
 
 func BenchmarkMarshal(b *testing.B) {