From ce83edbed157f50870d0a7ce06e1032fefba205f Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Fri, 21 May 2021 18:52:55 +0200
Subject: [PATCH] base: improve coverage

---
 pkg/base/body_test.go             | 20 +++++++++++++++++++
 pkg/base/header_test.go           | 24 ++++++++++++++++++++++
 pkg/base/interleavedframe.go      | 10 +++-------
 pkg/base/interleavedframe_test.go | 26 ++++++++++++++++++++++++
 pkg/base/request_test.go          | 33 +++++++++++++++++++++++++++++++
 pkg/base/response_test.go         | 33 +++++++++++++++++++++++++++++++
 6 files changed, 139 insertions(+), 7 deletions(-)

diff --git a/pkg/base/body_test.go b/pkg/base/body_test.go
index 65111a13..cd6a7d3c 100644
--- a/pkg/base/body_test.go
+++ b/pkg/base/body_test.go
@@ -3,11 +3,25 @@ package base
 import (
 	"bufio"
 	"bytes"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 )
 
+type limitedBuffer struct {
+	cap int
+	n   int
+}
+
+func (b *limitedBuffer) Write(p []byte) (int, error) {
+	b.n += len(p)
+	if b.n > b.cap {
+		return 0, fmt.Errorf("capacity reached")
+	}
+	return len(p), nil
+}
+
 var casesBody = []struct {
 	name string
 	h    Header
@@ -90,3 +104,9 @@ func TestBodyReadErrors(t *testing.T) {
 		})
 	}
 }
+
+func TestBodyWriteErrors(t *testing.T) {
+	bw := bufio.NewWriterSize(&limitedBuffer{cap: 3}, 1)
+	err := body([]byte("1234")).write(bw)
+	require.Equal(t, "capacity reached", err.Error())
+}
diff --git a/pkg/base/header_test.go b/pkg/base/header_test.go
index 25c8a054..4c62042d 100644
--- a/pkg/base/header_test.go
+++ b/pkg/base/header_test.go
@@ -183,3 +183,27 @@ func TestHeaderReadErrors(t *testing.T) {
 		})
 	}
 }
+
+func TestHeaderWriteErrors(t *testing.T) {
+	for _, ca := range []struct {
+		name string
+		cap  int
+	}{
+		{
+			"values",
+			3,
+		},
+		{
+			"final newline",
+			12,
+		},
+	} {
+		t.Run(ca.name, func(t *testing.T) {
+			bw := bufio.NewWriterSize(&limitedBuffer{cap: ca.cap}, 1)
+			err := Header{
+				"Value": HeaderValue{"key"},
+			}.write(bw)
+			require.Equal(t, "capacity reached", err.Error())
+		})
+	}
+}
diff --git a/pkg/base/interleavedframe.go b/pkg/base/interleavedframe.go
index bef6141d..8002eae9 100644
--- a/pkg/base/interleavedframe.go
+++ b/pkg/base/interleavedframe.go
@@ -116,14 +116,10 @@ func (f InterleavedFrame) Write(bw *bufio.Writer) error {
 		return uint8((f.TrackID * 2) + 1)
 	}()
 
-	_, err := bw.Write([]byte{0x24, channel})
-	if err != nil {
-		return err
-	}
+	buf := []byte{0x24, channel, 0x00, 0x00}
+	binary.BigEndian.PutUint16(buf[2:], uint16(len(f.Payload)))
 
-	buf := make([]byte, 2)
-	binary.BigEndian.PutUint16(buf, uint16(len(f.Payload)))
-	_, err = bw.Write(buf)
+	_, err := bw.Write(buf)
 	if err != nil {
 		return err
 	}
diff --git a/pkg/base/interleavedframe_test.go b/pkg/base/interleavedframe_test.go
index 6ad9c0a6..787f74ca 100644
--- a/pkg/base/interleavedframe_test.go
+++ b/pkg/base/interleavedframe_test.go
@@ -96,6 +96,32 @@ func TestInterleavedFrameReadErrors(t *testing.T) {
 	}
 }
 
+func TestInterleavedFrameWriteErrors(t *testing.T) {
+	for _, ca := range []struct {
+		name string
+		cap  int
+	}{
+		{
+			"header",
+			3,
+		},
+		{
+			"content",
+			6,
+		},
+	} {
+		t.Run(ca.name, func(t *testing.T) {
+			bw := bufio.NewWriterSize(&limitedBuffer{cap: ca.cap}, 1)
+			err := InterleavedFrame{
+				TrackID:    3,
+				StreamType: StreamTypeRTP,
+				Payload:    []byte{0x01, 0x02, 0x03, 0x04},
+			}.Write(bw)
+			require.Equal(t, "capacity reached", err.Error())
+		})
+	}
+}
+
 func TestReadInterleavedFrameOrRequest(t *testing.T) {
 	byts := []byte("DESCRIBE rtsp://example.com/media.mp4 RTSP/1.0\r\n" +
 		"Accept: application/sdp\r\n" +
diff --git a/pkg/base/request_test.go b/pkg/base/request_test.go
index ca2bfc5a..d96a6fda 100644
--- a/pkg/base/request_test.go
+++ b/pkg/base/request_test.go
@@ -231,6 +231,39 @@ func TestRequestReadErrors(t *testing.T) {
 	}
 }
 
+func TestRequestWriteErrors(t *testing.T) {
+	for _, ca := range []struct {
+		name string
+		cap  int
+	}{
+		{
+			"first line",
+			3,
+		},
+		{
+			"header",
+			20,
+		},
+		{
+			"body",
+			40,
+		},
+	} {
+		t.Run(ca.name, func(t *testing.T) {
+			bw := bufio.NewWriterSize(&limitedBuffer{cap: ca.cap}, 1)
+			err := Request{
+				Method: "ANNOUNCE",
+				URL:    MustParseURL("rtsp://example.com/media.mp4"),
+				Header: Header{
+					"CSeq": HeaderValue{"7"},
+				},
+				Body: []byte("abc"),
+			}.Write(bw)
+			require.Equal(t, "capacity reached", err.Error())
+		})
+	}
+}
+
 func TestRequestReadIgnoreFrames(t *testing.T) {
 	byts := []byte{0x24, 0x6, 0x0, 0x4, 0x1, 0x2, 0x3, 0x4}
 	byts = append(byts, []byte("OPTIONS rtsp://example.com/media.mp4 RTSP/1.0\r\n"+
diff --git a/pkg/base/response_test.go b/pkg/base/response_test.go
index f6bece06..e1fab8dd 100644
--- a/pkg/base/response_test.go
+++ b/pkg/base/response_test.go
@@ -188,6 +188,39 @@ func TestResponseReadErrors(t *testing.T) {
 	}
 }
 
+func TestResponseWriteErrors(t *testing.T) {
+	for _, ca := range []struct {
+		name string
+		cap  int
+	}{
+		{
+			"first line",
+			3,
+		},
+		{
+			"header",
+			20,
+		},
+		{
+			"body",
+			40,
+		},
+	} {
+		t.Run(ca.name, func(t *testing.T) {
+			bw := bufio.NewWriterSize(&limitedBuffer{cap: ca.cap}, 1)
+			err := Response{
+				StatusCode:    200,
+				StatusMessage: "OK",
+				Header: Header{
+					"CSeq": HeaderValue{"2"},
+				},
+				Body: []byte("abc"),
+			}.Write(bw)
+			require.Equal(t, "capacity reached", err.Error())
+		})
+	}
+}
+
 func TestResponseWriteAutoFillStatus(t *testing.T) {
 	res := &Response{
 		StatusCode: StatusMethodNotAllowed,