Fix the comments

14 files changed, 165 insertions, 95 deletions

diff --git a/go/cmd/gitlab-shell/main.go b/go/cmd/gitlab-shell/main.go
index 0b794e2..f7b6367 100644
--- a/go/cmd/gitlab-shell/main.go
+++ b/go/cmd/gitlab-shell/main.go
@@ -12,23 +12,23 @@ import (
 )
 
 var (
-	binDir   string
-	rootDir  string
-	reporter *reporting.Reporter
+	binDir     string
+	rootDir    string
+	readWriter *reporting.ReadWriter
 )
 
 func init() {
 	binDir = filepath.Dir(os.Args[0])
 	rootDir = filepath.Dir(binDir)
-	reporter = &reporting.Reporter{Out: os.Stdout, In: os.Stdin, ErrOut: os.Stderr}
+	readWriter = &reporting.ReadWriter{Out: os.Stdout, In: os.Stdin, ErrOut: os.Stderr}
 }
 
 // rubyExec will never return. It either replaces the current process with a
 // Ruby interpreter, or outputs an error and kills the process.
 func execRuby() {
 	cmd := &fallback.Command{}
-	if err := cmd.Execute(reporter); err != nil {
-		fmt.Fprintf(reporter.ErrOut, "Failed to exec: %v\n", err)
+	if err := cmd.Execute(readWriter); err != nil {
+		fmt.Fprintf(readWriter.ErrOut, "Failed to exec: %v\n", err)
 		os.Exit(1)
 	}
 }
@@ -38,7 +38,7 @@ func main() {
 	// warning as this isn't something we can sustain indefinitely
 	config, err := config.NewFromDir(rootDir)
 	if err != nil {
-		fmt.Fprintln(reporter.ErrOut, "Failed to read config, falling back to gitlab-shell-ruby")
+		fmt.Fprintln(readWriter.ErrOut, "Failed to read config, falling back to gitlab-shell-ruby")
 		execRuby()
 	}
 
@@ -46,14 +46,14 @@ func main() {
 	if err != nil {
 		// For now this could happen if `SSH_CONNECTION` is not set on
 		// the environment
-		fmt.Fprintf(reporter.ErrOut, "%v\n", err)
+		fmt.Fprintf(readWriter.ErrOut, "%v\n", err)
 		os.Exit(1)
 	}
 
 	// The command will write to STDOUT on execution or replace the current
 	// process in case of the `fallback.Command`
-	if err = cmd.Execute(reporter); err != nil {
-		fmt.Fprintf(reporter.ErrOut, "%v\n", err)
+	if err = cmd.Execute(readWriter); err != nil {
+		fmt.Fprintf(readWriter.ErrOut, "%v\n", err)
 		os.Exit(1)
 	}
 }
diff --git a/go/internal/command/command.go b/go/internal/command/command.go
index b0b7c61..a376d1d 100644
--- a/go/internal/command/command.go
+++ b/go/internal/command/command.go
@@ -5,12 +5,12 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/discover"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/fallback"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/reporting"
-	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/twofactorrecovery"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/twofactorrecover"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
 )
 
 type Command interface {
-	Execute(*reporting.Reporter) error
+	Execute(*reporting.ReadWriter) error
 }
 
 func New(arguments []string, config *config.Config) (Command, error) {
@@ -31,8 +31,8 @@ func buildCommand(args *commandargs.CommandArgs, config *config.Config) Command
 	switch args.CommandType {
 	case commandargs.Discover:
 		return &discover.Command{Config: config, Args: args}
-	case commandargs.TwoFactorRecovery:
-		return &twofactorrecovery.Command{Config: config, Args: args}
+	case commandargs.TwoFactorRecover:
+		return &twofactorrecover.Command{Config: config, Args: args}
 	}
 
 	return nil
diff --git a/go/internal/command/command_test.go b/go/internal/command/command_test.go
index 02fc0d0..42c5112 100644
--- a/go/internal/command/command_test.go
+++ b/go/internal/command/command_test.go
@@ -6,6 +6,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/discover"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/fallback"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/twofactorrecover"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/testhelper"
 )
@@ -44,6 +45,19 @@ func TestNew(t *testing.T) {
 			},
 			expectedType: &fallback.Command{},
 		},
+		{
+			desc:      "it returns a TwoFactorRecover command if the feature is enabled",
+			arguments: []string{},
+			config: &config.Config{
+				GitlabUrl: "http+unix://gitlab.socket",
+				Migration: config.MigrationConfig{Enabled: true, Features: []string{"2fa_recovery_codes"}},
+			},
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": "2fa_recovery_codes",
+			},
+			expectedType: &twofactorrecover.Command{},
+		},
 	}
 
 	for _, tc := range testCases {
diff --git a/go/internal/command/commandargs/command_args.go b/go/internal/command/commandargs/command_args.go
index 8e13c4b..e801889 100644
--- a/go/internal/command/commandargs/command_args.go
+++ b/go/internal/command/commandargs/command_args.go
@@ -9,8 +9,8 @@ import (
 type CommandType string
 
 const (
-	Discover          CommandType = "discover"
-	TwoFactorRecovery CommandType = "2fa_recovery_codes"
+	Discover         CommandType = "discover"
+	TwoFactorRecover CommandType = "2fa_recovery_codes"
 )
 
 var (
@@ -81,7 +81,7 @@ func (c *CommandArgs) parseCommand(commandString string) {
 		c.CommandType = Discover
 	}
 
-	if CommandType(commandString) == TwoFactorRecovery {
-		c.CommandType = TwoFactorRecovery
+	if CommandType(commandString) == TwoFactorRecover {
+		c.CommandType = TwoFactorRecover
 	}
 }
diff --git a/go/internal/command/discover/discover.go b/go/internal/command/discover/discover.go
index 8ad2868..6e80f21 100644
--- a/go/internal/command/discover/discover.go
+++ b/go/internal/command/discover/discover.go
@@ -14,16 +14,16 @@ type Command struct {
 	Args   *commandargs.CommandArgs
 }
 
-func (c *Command) Execute(reporter *reporting.Reporter) error {
+func (c *Command) Execute(readWriter *reporting.ReadWriter) error {
 	response, err := c.getUserInfo()
 	if err != nil {
 		return fmt.Errorf("Failed to get username: %v", err)
 	}
 
 	if response.IsAnonymous() {
-		fmt.Fprintf(reporter.Out, "Welcome to GitLab, Anonymous!\n")
+		fmt.Fprintf(readWriter.Out, "Welcome to GitLab, Anonymous!\n")
 	} else {
-		fmt.Fprintf(reporter.Out, "Welcome to GitLab, @%s!\n", response.Username)
+		fmt.Fprintf(readWriter.Out, "Welcome to GitLab, @%s!\n", response.Username)
 	}
 
 	return nil
diff --git a/go/internal/command/discover/discover_test.go b/go/internal/command/discover/discover_test.go
index ec6f931..673d848 100644
--- a/go/internal/command/discover/discover_test.go
+++ b/go/internal/command/discover/discover_test.go
@@ -82,7 +82,7 @@ func TestExecute(t *testing.T) {
 			cmd := &Command{Config: testConfig, Args: tc.arguments}
 			buffer := &bytes.Buffer{}
 
-			err := cmd.Execute(&reporting.Reporter{Out: buffer})
+			err := cmd.Execute(&reporting.ReadWriter{Out: buffer})
 
 			assert.NoError(t, err)
 			assert.Equal(t, tc.expectedOutput, buffer.String())
@@ -122,7 +122,7 @@ func TestFailingExecute(t *testing.T) {
 			cmd := &Command{Config: testConfig, Args: tc.arguments}
 			buffer := &bytes.Buffer{}
 
-			err := cmd.Execute(&reporting.Reporter{Out: buffer})
+			err := cmd.Execute(&reporting.ReadWriter{Out: buffer})
 
 			assert.Empty(t, buffer.String())
 			assert.EqualError(t, err, tc.expectedError)
diff --git a/go/internal/command/fallback/fallback.go b/go/internal/command/fallback/fallback.go
index a2c73ed..445cd71 100644
--- a/go/internal/command/fallback/fallback.go
+++ b/go/internal/command/fallback/fallback.go
@@ -14,7 +14,7 @@ var (
 	binDir = filepath.Dir(os.Args[0])
 )
 
-func (c *Command) Execute(_ *reporting.Reporter) error {
+func (c *Command) Execute(_ *reporting.ReadWriter) error {
 	rubyCmd := filepath.Join(binDir, "gitlab-shell-ruby")
 	execErr := syscall.Exec(rubyCmd, os.Args, os.Environ())
 	return execErr
diff --git a/go/internal/command/reporting/reporter.go b/go/internal/command/reporting/read_writer.go
index d1337b0..c0fb3da 100644
--- a/go/internal/command/reporting/reporter.go
+++ b/go/internal/command/reporting/read_writer.go
@@ -2,7 +2,7 @@ package reporting
 
 import "io"
 
-type Reporter struct {
+type ReadWriter struct {
 	Out    io.Writer
 	In     io.Reader
 	ErrOut io.Writer
diff --git a/go/internal/command/twofactorrecover/twofactorrecover.go b/go/internal/command/twofactorrecover/twofactorrecover.go
new file mode 100644
index 0000000..adfc511
--- /dev/null
+++ b/go/internal/command/twofactorrecover/twofactorrecover.go
@@ -0,0 +1,68 @@
+package twofactorrecover
+
+import (
+	"fmt"
+	"strings"
+
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/commandargs"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/reporting"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet/twofactorrecover"
+)
+
+type Command struct {
+	Config *config.Config
+	Args   *commandargs.CommandArgs
+}
+
+func (c *Command) Execute(readWriter *reporting.ReadWriter) error {
+	if c.canContinue(readWriter) {
+		c.displayRecoveryCodes(readWriter)
+	} else {
+		fmt.Fprintln(readWriter.Out, "\nNew recovery codes have *not* been generated. Existing codes will remain valid.")
+	}
+
+	return nil
+}
+
+func (c *Command) canContinue(readWriter *reporting.ReadWriter) bool {
+	question :=
+		"Are you sure you want to generate new two-factor recovery codes?\n" +
+			"Any existing recovery codes you saved will be invalidated. (yes/no)"
+	fmt.Fprintln(readWriter.Out, question)
+
+	var answer string
+	fmt.Fscanln(readWriter.In, &answer)
+
+	return (answer == "yes")
+}
+
+func (c *Command) displayRecoveryCodes(readWriter *reporting.ReadWriter) {
+	codes, err := c.getRecoveryCodes()
+
+	if err == nil {
+		messageWithCodes :=
+			"\nYour two-factor authentication recovery codes are:\n\n" +
+				strings.Join(codes, "\n") +
+				"\n\nDuring sign in, use one of the codes above when prompted for\n" +
+				"your two-factor code. Then, visit your Profile Settings and add\n" +
+				"a new device so you do not lose access to your account again.\n"
+		fmt.Fprint(readWriter.Out, messageWithCodes)
+	} else {
+		fmt.Fprintf(readWriter.Out, "\nAn error occurred while trying to generate new recovery codes.\n%v\n", err)
+	}
+}
+
+func (c *Command) getRecoveryCodes() ([]string, error) {
+	if c.Args.GitlabKeyId == "" {
+		return nil, fmt.Errorf("Failed to get key id")
+	}
+
+	client, err := twofactorrecover.NewClient(c.Config)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return client.GetRecoveryCodes(c.Args.GitlabKeyId)
+}
diff --git a/go/internal/command/twofactorrecovery/twofactorrecovery_test.go b/go/internal/command/twofactorrecover/twofactorrecover_test.go
index 6b988b0..28b8751 100644
--- a/go/internal/command/twofactorrecovery/twofactorrecovery_test.go
+++ b/go/internal/command/twofactorrecover/twofactorrecover_test.go
@@ -1,4 +1,4 @@
-package twofactorrecovery
+package twofactorrecover
 
 import (
 	"bytes"
@@ -92,7 +92,7 @@ func TestExecute(t *testing.T) {
 			arguments: &commandargs.CommandArgs{},
 			expectedOutput: "Are you sure you want to generate new two-factor recovery codes?\n" +
 				"Any existing recovery codes you saved will be invalidated. (yes/no)\n\n" +
-				"An error occurred while trying to generate new recovery codes.\n\n",
+				"An error occurred while trying to generate new recovery codes.\nFailed to get key id\n",
 		},
 	}
 
@@ -103,7 +103,7 @@ func TestExecute(t *testing.T) {
 
 			cmd := &Command{Config: testConfig, Args: tc.arguments}
 
-			err := cmd.Execute(&reporting.Reporter{Out: output, In: input})
+			err := cmd.Execute(&reporting.ReadWriter{Out: output, In: input})
 
 			assert.NoError(t, err)
 			assert.Equal(t, tc.expectedOutput, output.String())
diff --git a/go/internal/command/twofactorrecovery/twofactorrecovery.go b/go/internal/command/twofactorrecovery/twofactorrecovery.go
deleted file mode 100644
index 7504235..0000000
--- a/go/internal/command/twofactorrecovery/twofactorrecovery.go
+++ /dev/null
@@ -1,60 +0,0 @@
-package twofactorrecovery
-
-import (
-	"fmt"
-	"strings"
-
-	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/commandargs"
-	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/reporting"
-	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/twofactorrecovery/twofactorrecoveryclient"
-	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
-)
-
-type Command struct {
-	Config *config.Config
-	Args   *commandargs.CommandArgs
-}
-
-func (c *Command) Execute(reporter *reporting.Reporter) error {
-	if c.canContinue(reporter) {
-		c.displayRecoveryCodes(reporter)
-	} else {
-		fmt.Fprintln(reporter.Out, "\nNew recovery codes have *not* been generated. Existing codes will remain valid.")
-	}
-
-	return nil
-}
-
-func (c *Command) canContinue(reporter *reporting.Reporter) bool {
-	fmt.Fprintln(reporter.Out, "Are you sure you want to generate new two-factor recovery codes?")
-	fmt.Fprintln(reporter.Out, "Any existing recovery codes you saved will be invalidated. (yes/no)")
-
-	var answer string
-	fmt.Fscanln(reporter.In, &answer)
-
-	return (answer == "yes")
-}
-
-func (c *Command) displayRecoveryCodes(reporter *reporting.Reporter) {
-	codes, err := c.getRecoveryCodes()
-
-	if err == nil {
-		fmt.Fprint(reporter.Out, "\nYour two-factor authentication recovery codes are:\n\n")
-		fmt.Fprintln(reporter.Out, strings.Join(codes, "\n"))
-		fmt.Fprintln(reporter.Out, "\nDuring sign in, use one of the codes above when prompted for")
-		fmt.Fprintln(reporter.Out, "your two-factor code. Then, visit your Profile Settings and add")
-		fmt.Fprintln(reporter.Out, "a new device so you do not lose access to your account again.")
-	} else {
-		fmt.Fprintf(reporter.Out, "\nAn error occurred while trying to generate new recovery codes.\n%v\n", err)
-	}
-}
-
-func (c *Command) getRecoveryCodes() ([]string, error) {
-	client, err := twofactorrecoveryclient.GetClient(c.Config)
-
-	if err != nil {
-		return nil, err
-	}
-
-	return client.GetRecoveryCodes(c.Args.GitlabKeyId)
-}
diff --git a/go/internal/command/twofactorrecovery/twofactorrecoveryclient/twofactorrecoveryclient.go b/go/internal/gitlabnet/twofactorrecover/client.go
index e7ec709..bce0fcb 100644
--- a/go/internal/command/twofactorrecovery/twofactorrecoveryclient/twofactorrecoveryclient.go
+++ b/go/internal/gitlabnet/twofactorrecover/client.go
@@ -1,4 +1,4 @@
-package twofactorrecoveryclient
+package twofactorrecover
 
 import (
 	"encoding/json"
@@ -22,7 +22,7 @@ type Response struct {
 	Message       string   `json:"message"`
 }
 
-func GetClient(config *config.Config) (*Client, error) {
+func NewClient(config *config.Config) (*Client, error) {
 	client, err := gitlabnet.GetClient(config)
 	if err != nil {
 		return nil, fmt.Errorf("Error creating http client: %v", err)
@@ -42,10 +42,11 @@ func (c *Client) GetRecoveryCodes(gitlabKeyId string) ([]string, error) {
 		return nil, err
 	}
 
+	defer response.Body.Close()
 	parsedResponse, err := c.parseResponse(response)
 
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("Parsing failed")
 	}
 
 	if parsedResponse.Success {
@@ -56,7 +57,6 @@ func (c *Client) GetRecoveryCodes(gitlabKeyId string) ([]string, error) {
 }
 
 func (c *Client) parseResponse(resp *http.Response) (*Response, error) {
-	defer resp.Body.Close()
 	parsedResponse := &Response{}
 
 	if err := json.NewDecoder(resp.Body).Decode(parsedResponse); err != nil {
diff --git a/go/internal/command/twofactorrecovery/twofactorrecoveryclient/twofactorrecoveryclient_test.go b/go/internal/gitlabnet/twofactorrecover/client_test.go
index 6020baa..094b029 100644
--- a/go/internal/command/twofactorrecovery/twofactorrecoveryclient/twofactorrecoveryclient_test.go
+++ b/go/internal/gitlabnet/twofactorrecover/client_test.go
@@ -1,4 +1,4 @@
-package twofactorrecoveryclient
+package twofactorrecover
 
 import (
 	"encoding/json"
@@ -8,7 +8,9 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet/testserver"
 )
 
@@ -35,6 +37,16 @@ func init() {
 						"message": "missing user",
 					}
 					json.NewEncoder(w).Encode(body)
+				} else if r.URL.Query().Get("key_id") == "2" {
+					w.WriteHeader(http.StatusForbidden)
+					body := &gitlabnet.ErrorResponse{
+						Message: "Not allowed!",
+					}
+					json.NewEncoder(w).Encode(body)
+				} else if r.URL.Query().Get("key_id") == "3" {
+					w.Write([]byte("{ \"message\": \"broken json!\""))
+				} else if r.URL.Query().Get("key_id") == "4" {
+					w.WriteHeader(http.StatusForbidden)
 				} else {
 					fmt.Fprint(w, "null")
 				}
@@ -60,11 +72,47 @@ func TestMissingUser(t *testing.T) {
 	assert.Equal(t, "missing user", err.Error())
 }
 
+func TestErrorResponses(t *testing.T) {
+	client, cleanup := setup(t)
+	defer cleanup()
+
+	testCases := []struct {
+		desc          string
+		fakeId        string
+		expectedError string
+	}{
+		{
+			desc:          "A response with an error message",
+			fakeId:        "2",
+			expectedError: "Not allowed!",
+		},
+		{
+			desc:          "A response with bad JSON",
+			fakeId:        "3",
+			expectedError: "Parsing failed",
+		},
+		{
+			desc:          "An error response without message",
+			fakeId:        "4",
+			expectedError: "Internal API error (403)",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			resp, err := client.GetRecoveryCodes(tc.fakeId)
+
+			assert.EqualError(t, err, tc.expectedError)
+			assert.Nil(t, resp)
+		})
+	}
+}
+
 func setup(t *testing.T) (*Client, func()) {
 	cleanup, err := testserver.StartSocketHttpServer(requests)
 	require.NoError(t, err)
 
-	client, err := GetClient(testConfig)
+	client, err := NewClient(testConfig)
 	require.NoError(t, err)
 
 	return client, cleanup
diff --git a/spec/gitlab_shell_gitlab_shell_spec.rb b/spec/gitlab_shell_gitlab_shell_spec.rb
index c36be2e..6d6e172 100644
--- a/spec/gitlab_shell_gitlab_shell_spec.rb
+++ b/spec/gitlab_shell_gitlab_shell_spec.rb
@@ -117,7 +117,7 @@ describe 'bin/gitlab-shell' do
       write_config(
         "gitlab_url" => "http+unix://#{CGI.escape(tmp_socket_path)}",
         "migration" => { "enabled" => true,
-                        "features" => ["discover", "2fa_recovery_codes"] }
+                        "features" => ["discover"] }
       )
     end