Unit test sshd.handleEnv function

1 files changed, 44 insertions, 0 deletions

diff --git a/internal/sshd/session_test.go b/internal/sshd/session_test.go
new file mode 100644
index 0000000..152506e
--- /dev/null
+++ b/internal/sshd/session_test.go
@@ -0,0 +1,44 @@
+package sshd
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"golang.org/x/crypto/ssh"
+)
+
+func TestHandleEnv(t *testing.T) {
+	testCases := []struct {
+		desc                    string
+		payload                 []byte
+		expectedProtocolVersion string
+		expectedResult          bool
+	}{
+		{
+			desc:                    "invalid payload",
+			payload:                 []byte("invalid"),
+			expectedProtocolVersion: "1",
+			expectedResult:          false,
+		}, {
+			desc:                    "valid payload",
+			payload:                 ssh.Marshal(envRequest{Name: "GIT_PROTOCOL", Value: "2"}),
+			expectedProtocolVersion: "2",
+			expectedResult:          true,
+		}, {
+			desc:                    "valid payload with forbidden env var",
+			payload:                 ssh.Marshal(envRequest{Name: "GIT_PROTOCOL_ENV", Value: "2"}),
+			expectedProtocolVersion: "1",
+			expectedResult:          true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			s := &session{gitProtocolVersion: "1"}
+			r := &ssh.Request{Payload: tc.payload}
+
+			require.Equal(t, s.handleEnv(r), tc.expectedResult)
+			require.Equal(t, s.gitProtocolVersion, tc.expectedProtocolVersion)
+		})
+	}
+}