From 15f09836286fc219d1f7de4f7c8c06062fbc56aa Mon Sep 17 00:00:00 2001 From: aler9 <46489434+aler9@users.noreply.github.com> Date: Sat, 3 Apr 2021 11:02:06 +0200 Subject: [PATCH] add h264 utilities --- internal/clientrtmp/client.go | 24 +++-- internal/h264/annexb.go | 166 ++++++++++++++++++++++++++++++++++ internal/h264/annexb_test.go | 140 ++++++++++++++++++++++++++++ internal/h264/avcc.go | 55 +++++++++++ internal/h264/avcc_test.go | 85 +++++++++++++++++ internal/sourcertmp/source.go | 12 +-- 6 files changed, 467 insertions(+), 15 deletions(-) create mode 100644 internal/h264/annexb.go create mode 100644 internal/h264/annexb_test.go create mode 100644 internal/h264/avcc.go create mode 100644 internal/h264/avcc_test.go diff --git a/internal/clientrtmp/client.go b/internal/clientrtmp/client.go index f0e5d9b8..ce9240e8 100644 --- a/internal/clientrtmp/client.go +++ b/internal/clientrtmp/client.go @@ -18,10 +18,11 @@ import ( "github.com/aler9/gortsplib/pkg/rtpaac" "github.com/aler9/gortsplib/pkg/rtph264" "github.com/notedit/rtmp/av" - "github.com/notedit/rtmp/codec/h264" + rh264 "github.com/notedit/rtmp/codec/h264" "github.com/aler9/rtsp-simple-server/internal/client" "github.com/aler9/rtsp-simple-server/internal/externalcmd" + "github.com/aler9/rtsp-simple-server/internal/h264" "github.com/aler9/rtsp-simple-server/internal/logger" "github.com/aler9/rtsp-simple-server/internal/rtcpsenderset" "github.com/aler9/rtsp-simple-server/internal/rtmputils" @@ -254,7 +255,7 @@ func (c *Client) runRead() { c.conn.WriteMetadata(videoTrack, audioTrack) if videoTrack != nil { - codec := h264.Codec{ + codec := rh264.Codec{ SPS: map[int][]byte{ 0: h264SPS, }, @@ -349,14 +350,19 @@ func (c *Client) runRead() { // aggregate NALUs by PTS if nt.Timestamp != videoPTS { + data, err := h264.EncodeAVCC(videoBuf) + if err != nil { + return err + } + pkt := av.Packet{ Type: av.H264, - Data: h264.FillNALUsAVCC(videoBuf), + Data: data, Time: now.Sub(videoStartDTS), } c.conn.NetConn().SetWriteDeadline(time.Now().Add(c.writeTimeout)) - err := c.conn.WritePacket(pkt) + err = c.conn.WritePacket(pkt) if err != nil { return err } @@ -550,16 +556,16 @@ func (c *Client) runPublish() { return fmt.Errorf("ERR: received an H264 frame, but track is not set up") } - // decode from AVCC format - nalus, typ := h264.SplitNALUs(pkt.Data) - if typ != h264.NALU_AVCC { - return fmt.Errorf("invalid NALU format (%d)", typ) + nalus, err := h264.DecodeAVCC(pkt.Data) + if err != nil { + return err } + ts := pkt.Time + pkt.CTime var nts []*rtph264.NALUAndTimestamp for _, nt := range nalus { nts = append(nts, &rtph264.NALUAndTimestamp{ - Timestamp: pkt.Time + pkt.CTime, + Timestamp: ts, NALU: nt, }) } diff --git a/internal/h264/annexb.go b/internal/h264/annexb.go new file mode 100644 index 00000000..4da4e546 --- /dev/null +++ b/internal/h264/annexb.go @@ -0,0 +1,166 @@ +package h264 + +import ( + "fmt" +) + +func removeAntiCompetition(nalu []byte) []byte { + // 0x00 0x00 0x03 0x00 -> 0x00 0x00 0x00 + // 0x00 0x00 0x03 0x01 -> 0x00 0x00 0x01 + // 0x00 0x00 0x03 0x02 -> 0x00 0x00 0x02 + // 0x00 0x00 0x03 0x03 -> 0x00 0x00 0x03 + + var ret []byte + step := 0 + start := 0 + + for i, b := range nalu { + switch step { + case 0: + if b == 0 { + step++ + } + + case 1: + if b == 0 { + step++ + } else { + step = 0 + } + + case 2: + if b == 3 { + step++ + } else { + step = 0 + } + + case 3: + switch b { + case 3, 2, 1, 0: + ret = append(ret, nalu[start:i-3]...) + ret = append(ret, []byte{0x00, 0x00, b}...) + step = 0 + start = i + 1 + + default: + step = 0 + } + } + } + + ret = append(ret, nalu[start:]...) + + return ret +} + +func addAntiCompetition(dest []byte, nalu []byte) []byte { + step := 0 + start := 0 + + for i, b := range nalu { + switch step { + case 0: + if b == 0 { + step++ + } + + case 1: + if b == 0 { + step++ + } else { + step = 0 + } + + case 2: + switch b { + case 3, 2, 1, 0: + dest = append(dest, nalu[start:i-2]...) + dest = append(dest, []byte{0x00, 0x00, 0x03, b}...) + step = 0 + start = i + 1 + + default: + step = 0 + } + } + } + + dest = append(dest, nalu[start:]...) + + return dest +} + +// DecodeAnnexB decodes NALUs from the Annex-B code stream format. +func DecodeAnnexB(byts []byte) ([][]byte, error) { + bl := len(byts) + + // check initial delimiter + n := func() int { + if bl < 3 || byts[0] != 0x00 || byts[1] != 0x00 { + return -1 + } + + if byts[2] == 0x01 { + return 3 + } + + if bl < 4 || byts[2] != 0x00 || byts[3] != 0x01 { + return -1 + } + + return 4 + }() + if n < 0 { + return nil, fmt.Errorf("input doesn't start with a delimiter") + } + + var ret [][]byte + zeros := 0 + start := n + delimStart := 0 + + for i := n; i < bl; i++ { + switch byts[i] { + case 0: + if zeros == 0 { + delimStart = i + } + zeros++ + + case 1: + if zeros == 2 || zeros == 3 { + nalu := byts[start:delimStart] + if len(nalu) == 0 { + return nil, fmt.Errorf("empty NALU") + } + ret = append(ret, removeAntiCompetition(nalu)) + start = i + 1 + } + zeros = 0 + + default: + zeros = 0 + } + } + + nalu := byts[start:bl] + if len(nalu) == 0 { + return nil, fmt.Errorf("empty NALU") + } + ret = append(ret, removeAntiCompetition(nalu)) + + return ret, nil +} + +// EncodeAnnexB encodes NALUs into the Annex-B code stream format. +func EncodeAnnexB(nalus [][]byte) ([]byte, error) { + var ret []byte + + for _, nalu := range nalus { + ret = append(ret, []byte{0x00, 0x00, 0x00, 0x01}...) + ret = addAntiCompetition(ret, nalu) + } + + return ret, nil +} diff --git a/internal/h264/annexb_test.go b/internal/h264/annexb_test.go new file mode 100644 index 00000000..06f2b0dd --- /dev/null +++ b/internal/h264/annexb_test.go @@ -0,0 +1,140 @@ +package h264 + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +var annexBCases = []struct { + name string + encin []byte + encout []byte + dec [][]byte +}{ + { + "2 zeros, single", + []byte{0x00, 0x00, 0x01, 0xaa, 0xbb}, + []byte{0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb}, + [][]byte{ + {0xaa, 0xbb}, + }, + }, + { + "2 zeros, multiple", + []byte{ + 0x00, 0x00, 0x01, 0xaa, 0xbb, 0x00, 0x00, 0x01, + 0xcc, 0xdd, 0x00, 0x00, 0x01, 0xee, 0xff, + }, + []byte{ + 0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb, 0x00, 0x00, + 0x00, 0x01, 0xcc, 0xdd, 0x00, 0x00, 0x00, 0x01, + 0xee, 0xff, + }, + [][]byte{ + {0xaa, 0xbb}, + {0xcc, 0xdd}, + {0xee, 0xff}, + }, + }, + { + "3 zeros, single", + []byte{0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb}, + []byte{0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb}, + [][]byte{ + {0xaa, 0xbb}, + }, + }, + { + "3 zeros, multiple", + []byte{ + 0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb, 0x00, 0x00, + 0x00, 0x01, 0xcc, 0xdd, 0x00, 0x00, 0x00, 0x01, + 0xee, 0xff, + }, + []byte{ + 0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb, 0x00, 0x00, + 0x00, 0x01, 0xcc, 0xdd, 0x00, 0x00, 0x00, 0x01, + 0xee, 0xff, + }, + [][]byte{ + {0xaa, 0xbb}, + {0xcc, 0xdd}, + {0xee, 0xff}, + }, + }, + { + "anti-competition", + []byte{ + 0x00, 0x00, 0x00, 0x01, + 0x00, 0x00, 0x03, 0x00, + 0x00, 0x00, 0x03, 0x01, + 0x00, 0x00, 0x03, 0x02, + 0x00, 0x00, 0x03, 0x03, + }, + []byte{ + 0x00, 0x00, 0x00, 0x01, + 0x00, 0x00, 0x03, 0x00, + 0x00, 0x00, 0x03, 0x01, + 0x00, 0x00, 0x03, 0x02, + 0x00, 0x00, 0x03, 0x03, + }, + [][]byte{ + { + 0x00, 0x00, 0x00, + 0x00, 0x00, 0x01, + 0x00, 0x00, 0x02, + 0x00, 0x00, 0x03, + }, + }, + }, +} + +func TestAnnexBDecode(t *testing.T) { + for _, ca := range annexBCases { + t.Run(ca.name, func(t *testing.T) { + dec, err := DecodeAnnexB(ca.encin) + require.NoError(t, err) + require.Equal(t, ca.dec, dec) + }) + } +} + +func TestAnnexBEncode(t *testing.T) { + for _, ca := range annexBCases { + t.Run(ca.name, func(t *testing.T) { + enc, err := EncodeAnnexB(ca.dec) + require.NoError(t, err) + require.Equal(t, ca.encout, enc) + }) + } +} + +func TestAnnexBDecodeError(t *testing.T) { + for _, ca := range []struct { + name string + enc []byte + }{ + { + "empty", + []byte{}, + }, + { + "missing initial delimiter", + []byte{0xaa, 0xbb}, + }, + { + "empty initial", + []byte{0x00, 0x00, 0x01}, + }, + { + "empty 2nd", + []byte{0x00, 0x00, 0x01, 0xaa, 0x00, 0x00, 0x01}, + }, + } { + t.Run(ca.name, func(t *testing.T) { + _, err := DecodeAnnexB(ca.enc) + require.Error(t, err) + }) + } +} diff --git a/internal/h264/avcc.go b/internal/h264/avcc.go new file mode 100644 index 00000000..74be0446 --- /dev/null +++ b/internal/h264/avcc.go @@ -0,0 +1,55 @@ +package h264 + +import ( + "encoding/binary" + "fmt" +) + +// DecodeAVCC encodes NALUs from the AVCC code stream format. +func DecodeAVCC(byts []byte) ([][]byte, error) { + var ret [][]byte + + for len(byts) > 0 { + if len(byts) < 4 { + return nil, fmt.Errorf("invalid length") + } + + le := binary.BigEndian.Uint32(byts) + byts = byts[4:] + + if len(byts) < int(le) { + return nil, fmt.Errorf("invalid length") + } + + ret = append(ret, byts[:le]) + byts = byts[le:] + } + + if len(ret) == 0 { + return nil, fmt.Errorf("no NALUs decoded") + } + + return ret, nil +} + +// EncodeAVCC encodes NALUs into the AVCC code stream format. +func EncodeAVCC(nalus [][]byte) ([]byte, error) { + le := 0 + for _, nalu := range nalus { + le += 4 + len(nalu) + } + + ret := make([]byte, le) + pos := 0 + + for _, nalu := range nalus { + ln := len(nalu) + binary.BigEndian.PutUint32(ret[pos:], uint32(ln)) + pos += 4 + + copy(ret[pos:], nalu) + pos += ln + } + + return ret, nil +} diff --git a/internal/h264/avcc_test.go b/internal/h264/avcc_test.go new file mode 100644 index 00000000..401d0e34 --- /dev/null +++ b/internal/h264/avcc_test.go @@ -0,0 +1,85 @@ +package h264 + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +var casesAVCC = []struct { + name string + enc []byte + dec [][]byte +}{ + { + "single", + []byte{ + 0x00, 0x00, 0x00, 0x03, + 0xaa, 0xbb, 0xcc, + }, + [][]byte{ + {0xaa, 0xbb, 0xcc}, + }, + }, + { + "multiple", + []byte{ + 0x00, 0x00, 0x00, 0x02, + 0xaa, 0xbb, + 0x00, 0x00, 0x00, 0x02, + 0xcc, 0xdd, + 0x00, 0x00, 0x00, 0x02, + 0xee, 0xff, + }, + [][]byte{ + {0xaa, 0xbb}, + {0xcc, 0xdd}, + {0xee, 0xff}, + }, + }, +} + +func TestAVCCDecode(t *testing.T) { + for _, ca := range casesAVCC { + t.Run(ca.name, func(t *testing.T) { + dec, err := DecodeAVCC(ca.enc) + require.NoError(t, err) + require.Equal(t, ca.dec, dec) + }) + } +} + +func TestAVCCEncode(t *testing.T) { + for _, ca := range casesAVCC { + t.Run(ca.name, func(t *testing.T) { + enc, err := EncodeAVCC(ca.dec) + require.NoError(t, err) + require.Equal(t, ca.enc, enc) + }) + } +} + +func TestAVCCDecodeError(t *testing.T) { + for _, ca := range []struct { + name string + enc []byte + }{ + { + "empty", + []byte{}, + }, + { + "invalid length", + []byte{0x01}, + }, + { + "invalid length", + []byte{0x00, 0x00, 0x00, 0x03}, + }, + } { + t.Run(ca.name, func(t *testing.T) { + _, err := DecodeAVCC(ca.enc) + require.Error(t, err) + }) + } +} diff --git a/internal/sourcertmp/source.go b/internal/sourcertmp/source.go index 5c058b86..811bf6ae 100644 --- a/internal/sourcertmp/source.go +++ b/internal/sourcertmp/source.go @@ -11,9 +11,9 @@ import ( "github.com/aler9/gortsplib/pkg/rtpaac" "github.com/aler9/gortsplib/pkg/rtph264" "github.com/notedit/rtmp/av" - "github.com/notedit/rtmp/codec/h264" "github.com/notedit/rtmp/format/rtmp" + "github.com/aler9/rtsp-simple-server/internal/h264" "github.com/aler9/rtsp-simple-server/internal/logger" "github.com/aler9/rtsp-simple-server/internal/rtcpsenderset" "github.com/aler9/rtsp-simple-server/internal/rtmputils" @@ -214,16 +214,16 @@ func (s *Source) runInner() bool { return fmt.Errorf("ERR: received an H264 frame, but track is not set up") } - // decode from AVCC format - nalus, typ := h264.SplitNALUs(pkt.Data) - if typ != h264.NALU_AVCC { - return fmt.Errorf("invalid NALU format (%d)", typ) + nalus, err := h264.DecodeAVCC(pkt.Data) + if err != nil { + return err } + ts := pkt.Time + pkt.CTime var nts []*rtph264.NALUAndTimestamp for _, nt := range nalus { nts = append(nts, &rtph264.NALUAndTimestamp{ - Timestamp: pkt.Time + pkt.CTime, + Timestamp: ts, NALU: nt, }) }