Unit test wrong channel type

1 files changed, 34 insertions, 4 deletions

diff --git a/internal/sshd/connection_test.go b/internal/sshd/connection_test.go
index 03e9209..f48750e 100644
--- a/internal/sshd/connection_test.go
+++ b/internal/sshd/connection_test.go
@@ -8,9 +8,15 @@ import (
 	"golang.org/x/crypto/ssh"
 )
 
+type rejectCall struct {
+	reason ssh.RejectionReason
+	message string
+}
+
 type fakeNewChannel struct {
 	channelType string
 	extraData   []byte
+	rejectCh chan rejectCall
 }
 
 func (f *fakeNewChannel) Accept() (ssh.Channel, <-chan *ssh.Request, error) {
@@ -18,6 +24,8 @@ func (f *fakeNewChannel) Accept() (ssh.Channel, <-chan *ssh.Request, error) {
 }
 
 func (f *fakeNewChannel) Reject(reason ssh.RejectionReason, message string) error {
+	f.rejectCh <- rejectCall{reason: reason, message: message}
+
 	return nil
 }
 
@@ -29,14 +37,20 @@ func (f *fakeNewChannel) ExtraData() []byte {
 	return f.extraData
 }
 
-func TestPanicDuringSessionIsRecovered(t *testing.T) {
-	numSessions := 0
-	conn := newConnection(1, "127.0.0.1:50000")
+func setup(sessionsNum int64, newChannel *fakeNewChannel) (*connection, chan ssh.NewChannel) {
+	conn := newConnection(sessionsNum, "127.0.0.1:50000")
 
-	newChannel := &fakeNewChannel{channelType: "session"}
 	chans := make(chan ssh.NewChannel, 1)
 	chans <- newChannel
 
+	return conn, chans
+}
+
+func TestPanicDuringSessionIsRecovered(t *testing.T) {
+	newChannel := &fakeNewChannel{channelType: "session"}
+	conn, chans := setup(1, newChannel)
+
+	numSessions := 0
 	require.NotPanics(t, func() {
 		conn.handle(context.Background(), chans, func(context.Context, ssh.Channel, <-chan *ssh.Request) {
 			numSessions += 1
@@ -47,3 +61,19 @@ func TestPanicDuringSessionIsRecovered(t *testing.T) {
 
 	require.Equal(t, numSessions, 1)
 }
+
+func TestUnknownChannelType(t *testing.T) {
+	rejectCh := make(chan rejectCall, 1)
+	newChannel := &fakeNewChannel{channelType: "unknown session", rejectCh: rejectCh}
+	conn, chans := setup(1, newChannel)
+
+	go func() {
+		conn.handle(context.Background(), chans, nil)
+	}()
+
+	rejectionData := <-rejectCh
+	close(rejectCh)
+
+	expectedRejection := rejectCall{reason: ssh.UnknownChannelType, message: "unknown channel type"}
+	require.Equal(t, expectedRejection, rejectionData)
+}