From c15894dc124baf1f448c8f53367997c41badea2d Mon Sep 17 00:00:00 2001
From: Saxon <saxon@ausocean.org>
Date: Wed, 1 May 2019 16:45:22 +0930
Subject: [PATCH] codec/h265: completed lex_test.go and fixed bugs in lex.go

---
 codec/h265/lex.go      |   8 +-
 codec/h265/lex_test.go | 196 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 200 insertions(+), 4 deletions(-)

diff --git a/codec/h265/lex.go b/codec/h265/lex.go
index 61a4dab7..f9197af0 100644
--- a/codec/h265/lex.go
+++ b/codec/h265/lex.go
@@ -67,7 +67,7 @@ func (l *Lexer) Lex(dst io.Writer, src io.Reader, delay time.Duration) error {
 		switch err {
 		case nil:
 		case io.EOF:
-			continue
+			return nil
 		default:
 			return fmt.Errorf("source read error: %v\n", err)
 		}
@@ -78,7 +78,7 @@ func (l *Lexer) Lex(dst io.Writer, src io.Reader, delay time.Duration) error {
 			return fmt.Errorf("could not get rtp payload, failed with err: %v\n", err)
 		}
 
-		nalType := (buf[0] >> 1) & 0x3f
+		nalType := (payload[0] >> 1) & 0x3f
 
 		// If not currently fragmented then we ignore current write
 		if l.fragmented && nalType != typeFragmentation {
@@ -143,8 +143,8 @@ func (l *Lexer) handleFragmentation(d []byte) {
 		_d = d[5:]
 	}
 
-	s := d[2]&0x80 == 1
-	e := d[2]&0x40 == 1
+	s := d[2]&0x80 != 0
+	e := d[2]&0x40 != 0
 	switch {
 	case s && !e:
 		l.fragmented = true
diff --git a/codec/h265/lex_test.go b/codec/h265/lex_test.go
index 6ad208b7..2ae042d4 100644
--- a/codec/h265/lex_test.go
+++ b/codec/h265/lex_test.go
@@ -26,3 +26,199 @@ LICENSE
 */
 
 package h265
+
+import (
+	"io"
+	"testing"
+
+	"bitbucket.org/ausocean/av/protocol/rtp"
+)
+
+type rtpReader struct {
+	pkts []rtp.Pkt
+	idx  int
+}
+
+func (r *rtpReader) Read(p []byte) (int, error) {
+	if r.idx == len(r.pkts) {
+		return 0, io.EOF
+	}
+	b := r.pkts[r.idx].Bytes(nil)
+	copy(p, r.pkts[r.idx].Bytes(nil))
+	r.idx++
+	return len(b), nil
+}
+
+type destination [][]byte
+
+func (d *destination) Write(p []byte) (int, error) {
+	t := make([]byte, len(p))
+	copy(t, p)
+	*d = append([][]byte(*d), t)
+	return len(p), nil
+}
+
+func TestLex(t *testing.T) {
+	const rtpVer = 2
+
+	packets := []rtp.Pkt{
+		// Singla NAL
+		rtp.Pkt{
+			V: rtpVer,
+			Payload: []byte{
+				// NAL header (type=32 VPS)
+				0x40, 0x00,
+				// NAL data
+				0x01, 0x02, 0x03, 0x04,
+			},
+		},
+
+		// Fragmentation (start packet)
+		rtp.Pkt{
+			V: rtpVer,
+			Payload: []byte{
+				// NAL header (type 49)
+				0x62, 0x00,
+				// FU header
+				0x80,
+				// FU Payload
+				0x01, 0x02, 0x03,
+			},
+		},
+
+		// Fragmentation (middle packet)
+		rtp.Pkt{
+			V: rtpVer,
+			Payload: []byte{
+				// NAL header (type 49)
+				0x62, 0x00,
+				// FU header
+				0x00,
+				// FU Payload
+				0x04, 0x05, 0x06,
+			},
+		},
+
+		// Fragmentation (end packet)
+		rtp.Pkt{
+			V: rtpVer,
+			Payload: []byte{
+				// NAL header (type 49)
+				0x62, 0x00,
+				// FU header
+				0x40,
+				// FU Payload
+				0x07, 0x08, 0x09,
+			},
+		},
+
+		// Aggregation.
+		// This is the last bit of data for the access unit, so set M true
+		rtp.Pkt{
+			V: rtpVer,
+			M: true,
+			Payload: []byte{
+				// NAL header (type 49)
+				0x60, 0x00,
+				// NAL1 Size
+				0x00, 0x04,
+				// NAL1 data
+				0x01, 0x02, 0x03, 0x04,
+				// NAL2 Size
+				0x00, 0x04,
+				// NAL2 data
+				0x01, 0x02, 0x03, 0x04,
+			},
+		},
+
+		// Singla NAL
+		rtp.Pkt{
+			V: rtpVer,
+			Payload: []byte{
+				// NAL header (type=32 VPS)
+				0x40, 0x00,
+				// NAL data
+				0x01, 0x02, 0x03, 0x04,
+			},
+		},
+
+		// Singla NAL
+		rtp.Pkt{
+			V: rtpVer,
+			M: true,
+			Payload: []byte{
+				// NAL header (type=32 VPS)
+				0x40, 0x00,
+				// NAL data
+				0x01, 0x02, 0x03, 0x04,
+			},
+		},
+	}
+
+	expect := [][]byte{
+		{
+			// NAL 1
+			// Start code
+			0x00, 0x00, 0x00, 0x01,
+			// NAL header (type=32 VPS)
+			0x40, 0x00,
+			// NAL data
+			0x01, 0x02, 0x03, 0x04,
+
+			// NAL 2
+			// Start code
+			0x00, 0x00, 0x00, 0x01,
+			// FU Payload
+			0x01, 0x02, 0x03,
+			// FU Payload
+			0x04, 0x05, 0x06,
+			// FU Payload
+			0x07, 0x08, 0x09,
+
+			// NAL 3
+			// Start code
+			0x00, 0x00, 0x00, 0x01,
+			// NAL data
+			0x01, 0x02, 0x03, 0x04,
+
+			// NAL 4
+			// Start code
+			0x00, 0x00, 0x00, 0x01,
+			// NAL2 data
+			0x01, 0x02, 0x03, 0x04,
+		},
+		{
+			// NAL 1
+			//Start code
+			0x00, 0x00, 0x00, 0x01,
+			// NAL header (type=32 VPS)
+			0x40, 0x00,
+			// Data
+			0x01, 0x02, 0x03, 0x04,
+
+			// NAL 2
+			//Start code
+			0x00, 0x00, 0x00, 0x01,
+			// NAL header (type=32 VPS)
+			0x40, 0x00,
+			// Data
+			0x01, 0x02, 0x03, 0x04,
+		},
+	}
+
+	r := &rtpReader{pkts: packets}
+	d := &destination{}
+	l := Lexer{UsingDON: false}
+	err := l.Lex(d, r, 0)
+	if err != nil {
+		t.Fatalf("error lexing: %v\n", err)
+	}
+
+	for i, au := range expect {
+		for j, b := range au {
+			if b != [][]byte(*d)[i][j] {
+				t.Fatalf("did not get expected data.\nGot: %v\nWant: %v\n", d, expect)
+			}
+		}
+	}
+}