From 954d46dce5846c94c2f3cfd07206acacc5208d19 Mon Sep 17 00:00:00 2001
From: Edmundo Gonzalez <51725820+edgonmsft@users.noreply.github.com>
Date: Mon, 30 Aug 2021 17:30:28 +0000
Subject: [PATCH] Changes from comments on pr.

---
 internal/codespaces/ssh.go    | 8 ++------
 internal/codespaces/states.go | 7 ++++++-
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/internal/codespaces/ssh.go b/internal/codespaces/ssh.go
index a1abcf381..1c3807742 100644
--- a/internal/codespaces/ssh.go
+++ b/internal/codespaces/ssh.go
@@ -15,9 +15,6 @@ import (
 	"github.com/github/go-liveshare"
 )
 
-// MakeSSHTunnel This function initializes the liveshare tunnel
-// Creates the tunnel from a local port to a remote port.
-// Returns the local port that was used, the channel and the error if any.
 func MakeSSHTunnel(ctx context.Context, lsclient *liveshare.Client, localSSHPort int, remoteSSHPort int) (int, <-chan error, error) {
 	tunnelClosed := make(chan error)
 
@@ -48,9 +45,8 @@ func MakeSSHTunnel(ctx context.Context, lsclient *liveshare.Client, localSSHPort
 	return port, tunnelClosed, nil
 }
 
-// StartSSHServer starts and installs the SSH server in the codespace
-// returns the remote port where it is running, the user to use to login
-// or an error if something failed.
+// StartSSHServer installs (if necessary) and starts the SSH in the codespace.
+// It returns the remote port where it is running, the user to log in with, or an error if something failed.
 func StartSSHServer(ctx context.Context, client *liveshare.Client) (serverPort int, user string, err error) {
 	sshServer, err := liveshare.NewSSHServer(client)
 	if err != nil {
diff --git a/internal/codespaces/states.go b/internal/codespaces/states.go
index 7b98b12dd..d2aa389ef 100644
--- a/internal/codespaces/states.go
+++ b/internal/codespaces/states.go
@@ -45,7 +45,12 @@ func PollPostCreateStates(ctx context.Context, log logger, apiClient *api.API, u
 		return fmt.Errorf("connect to liveshare: %v", err)
 	}
 
-	tunnelPort, connClosed, err := MakeSSHTunnel(ctx, lsclient, 0, 2222)
+	remoteSSHServerPort, _, err := StartSSHServer(ctx, lsclient)
+	if err != nil {
+		return fmt.Errorf("error getting ssh server details: %v", err)
+	}
+
+	tunnelPort, connClosed, err := MakeSSHTunnel(ctx, lsclient, 0, remoteSSHServerPort)
 	if err != nil {
 		return fmt.Errorf("make ssh tunnel: %v", err)
 	}