Merge branch '506-twofactorverify-command-to-support-push-notification' into 'main'

Implement Push Auth support for 2FA verification

Closes #506

See merge request gitlab-org/gitlab-shell!454

internal/command/twofactorverify/twofactorverify.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"time"
 
 	"gitlab.com/gitlab-org/labkit/log"
 
@@ -13,6 +14,11 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/gitlabnet/twofactorverify"
 )
 
+const (
+	timeout = 30 * time.Second
+	prompt  = "OTP: "
+)
+
 type Command struct {
 	Config     *config.Config
 	Args       *commandargs.Shell
@@ -20,25 +26,51 @@ type Command struct {
 }
 
 func (c *Command) Execute(ctx context.Context) error {
-	ctxlog := log.ContextLogger(ctx)
-	ctxlog.Info("twofactorverify: execute: waiting for user input")
-	otp := c.getOTP(ctx)
-
-	ctxlog.Info("twofactorverify: execute: verifying entered OTP")
-	err := c.verifyOTP(ctx, otp)
+	client, err := twofactorverify.NewClient(c.Config)
 	if err != nil {
-		ctxlog.WithError(err).Error("twofactorverify: execute: OTP verification failed")
 		return err
 	}
 
-	ctxlog.WithError(err).Info("twofactorverify: execute: OTP verified")
-	return nil
-}
+	ctx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
 
-func (c *Command) getOTP(ctx context.Context) string {
-	prompt := "OTP: "
 	fmt.Fprint(c.ReadWriter.Out, prompt)
 
+	resultCh := make(chan string)
+	go func() {
+		err := client.PushAuth(ctx, c.Args)
+		if err == nil {
+			resultCh <- "OTP has been validated by Push Authentication. Git operations are now allowed."
+		}
+	}()
+
+	go func() {
+		answer, err := c.getOTP(ctx)
+		if err != nil {
+			resultCh <- formatErr(err)
+		}
+
+		if err := client.VerifyOTP(ctx, c.Args, answer); err != nil {
+			resultCh <- formatErr(err)
+		} else {
+			resultCh <- "OTP validation successful. Git operations are now allowed."
+		}
+	}()
+
+	var message string
+	select {
+	case message = <-resultCh:
+	case <-ctx.Done():
+		message = formatErr(ctx.Err())
+	}
+
+	log.WithContextFields(ctx, log.Fields{"message": message}).Info("Two factor verify command finished")
+	fmt.Fprintf(c.ReadWriter.Out, "\n%v\n", message)
+
+	return nil
+}
+
+func (c *Command) getOTP(ctx context.Context) (string, error) {
 	var answer string
 	otpLength := int64(64)
 	reader := io.LimitReader(c.ReadWriter.In, otpLength)
@@ -46,21 +78,13 @@ func (c *Command) getOTP(ctx context.Context) string {
 		log.ContextLogger(ctx).WithError(err).Debug("twofactorverify: getOTP: Failed to get user input")
 	}
 
-	return answer
-}
-
-func (c *Command) verifyOTP(ctx context.Context, otp string) error {
-	client, err := twofactorverify.NewClient(c.Config)
-	if err != nil {
-		return err
+	if answer == "" {
+		return "", fmt.Errorf("OTP cannot be blank.")
 	}
 
-	err = client.VerifyOTP(ctx, c.Args, otp)
-	if err == nil {
-		fmt.Fprint(c.ReadWriter.Out, "\nOTP validation successful. Git operations are now allowed.\n")
-	} else {
-		fmt.Fprintf(c.ReadWriter.Out, "\nOTP validation failed.\n%v\n", err)
-	}
+	return answer, nil
+}
 
-	return nil
+func formatErr(err error) string {
+	return fmt.Sprintf("OTP validation failed: %v", err)
 }

internal/command/twofactorverify/twofactorverify_test.go

@@ -17,10 +17,20 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/gitlabnet/twofactorverify"
 )
 
+type blockingReader struct{}
+
+func (*blockingReader) Read([]byte) (int, error) {
+	waitInfinitely := make(chan struct{})
+	<-waitInfinitely
+
+	return 0, nil
+}
+
 func setup(t *testing.T) []testserver.TestRequestHandler {
+	waitInfinitely := make(chan struct{})
 	requests := []testserver.TestRequestHandler{
 		{
-			Path: "/api/v4/internal/two_factor_otp_check",
+			Path: "/api/v4/internal/two_factor_manual_otp_check",
 			Handler: func(w http.ResponseWriter, r *http.Request) {
 				b, err := io.ReadAll(r.Body)
 				defer r.Body.Close()
@@ -31,11 +41,13 @@ func setup(t *testing.T) []testserver.TestRequestHandler {
 				require.NoError(t, json.Unmarshal(b, &requestBody))
 
 				switch requestBody.KeyId {
-				case "1":
+				case "verify_via_otp", "verify_via_otp_with_push_error":
 					body := map[string]interface{}{
 						"success": true,
 					}
 					json.NewEncoder(w).Encode(body)
+				case "wait_infinitely":
+					<-waitInfinitely
 				case "error":
 					body := map[string]interface{}{
 						"success": false,
@@ -47,15 +59,36 @@ func setup(t *testing.T) []testserver.TestRequestHandler {
 				}
 			},
 		},
+		{
+			Path: "/api/v4/internal/two_factor_push_otp_check",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				b, err := io.ReadAll(r.Body)
+				defer r.Body.Close()
+
+				require.NoError(t, err)
+
+				var requestBody *twofactorverify.RequestBody
+				require.NoError(t, json.Unmarshal(b, &requestBody))
+
+				switch requestBody.KeyId {
+				case "verify_via_push":
+					body := map[string]interface{}{
+						"success": true,
+					}
+					json.NewEncoder(w).Encode(body)
+				case "verify_via_otp_with_push_error":
+					w.WriteHeader(http.StatusInternalServerError)
+				default:
+					<-waitInfinitely
+				}
+			},
+		},
 	}
 
 	return requests
 }
 
-const (
-	question    = "OTP: \n"
-	errorHeader = "OTP validation failed.\n"
-)
+const errorHeader = "OTP validation failed: "
 
 func TestExecute(t *testing.T) {
 	requests := setup(t)
@@ -65,46 +98,61 @@ func TestExecute(t *testing.T) {
 	testCases := []struct {
 		desc           string
 		arguments      *commandargs.Shell
-		answer         string
+		input          io.Reader
 		expectedOutput string
 	}{
 		{
-			desc:      "With a known key id",
-			arguments: &commandargs.Shell{GitlabKeyId: "1"},
-			answer:    "123456\n",
-			expectedOutput: question +
-				"OTP validation successful. Git operations are now allowed.\n",
+			desc:           "Verify via OTP",
+			arguments:      &commandargs.Shell{GitlabKeyId: "verify_via_otp"},
+			expectedOutput: "OTP validation successful. Git operations are now allowed.\n",
+		},
+		{
+			desc:           "Verify via OTP",
+			arguments:      &commandargs.Shell{GitlabKeyId: "verify_via_otp_with_push_error"},
+			expectedOutput: "OTP validation successful. Git operations are now allowed.\n",
+		},
+		{
+			desc:           "Verify via push authentication",
+			arguments:      &commandargs.Shell{GitlabKeyId: "verify_via_push"},
+			input:          &blockingReader{},
+			expectedOutput: "OTP has been validated by Push Authentication. Git operations are now allowed.\n",
+		},
+		{
+			desc:           "With an empty OTP",
+			arguments:      &commandargs.Shell{GitlabKeyId: "verify_via_otp"},
+			input:          bytes.NewBufferString("\n"),
+			expectedOutput: errorHeader + "OTP cannot be blank.\n",
 		},
 		{
 			desc:           "With bad response",
 			arguments:      &commandargs.Shell{GitlabKeyId: "-1"},
-			answer:         "123456\n",
-			expectedOutput: question + errorHeader + "Parsing failed\n",
+			expectedOutput: errorHeader + "Parsing failed\n",
 		},
 		{
 			desc:           "With API returns an error",
 			arguments:      &commandargs.Shell{GitlabKeyId: "error"},
-			answer:         "yes\n",
-			expectedOutput: question + errorHeader + "error message\n",
+			expectedOutput: errorHeader + "error message\n",
 		},
 		{
 			desc:           "With API fails",
 			arguments:      &commandargs.Shell{GitlabKeyId: "broken"},
-			answer:         "yes\n",
-			expectedOutput: question + errorHeader + "Internal API error (500)\n",
+			expectedOutput: errorHeader + "Internal API error (500)\n",
 		},
 		{
 			desc:           "With missing arguments",
 			arguments:      &commandargs.Shell{},
-			answer:         "yes\n",
-			expectedOutput: question + errorHeader + "who='' is invalid\n",
+			expectedOutput: errorHeader + "who='' is invalid\n",
 		},
 	}
 
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
 			output := &bytes.Buffer{}
-			input := bytes.NewBufferString(tc.answer)
+
+			input := tc.input
+			if input == nil {
+				input = bytes.NewBufferString("123456\n")
+			}
 
 			cmd := &Command{
 				Config:     &config.Config{GitlabUrl: url},
@@ -115,7 +163,29 @@ func TestExecute(t *testing.T) {
 			err := cmd.Execute(context.Background())
 
 			require.NoError(t, err)
-			require.Equal(t, tc.expectedOutput, output.String())
+			require.Equal(t, prompt+"\n"+tc.expectedOutput, output.String())
 		})
 	}
 }
+
+func TestCanceledContext(t *testing.T) {
+	requests := setup(t)
+
+	output := &bytes.Buffer{}
+
+	url := testserver.StartSocketHttpServer(t, requests)
+	cmd := &Command{
+		Config:     &config.Config{GitlabUrl: url},
+		Args:       &commandargs.Shell{GitlabKeyId: "wait_infinitely"},
+		ReadWriter: &readwriter.ReadWriter{Out: output, In: &bytes.Buffer{}},
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	errCh := make(chan error)
+	go func() { errCh <- cmd.Execute(ctx) }()
+	cancel()
+
+	require.NoError(t, <-errCh)
+	require.Equal(t, prompt+"\n"+errorHeader+"context canceled\n", output.String())
+}

internal/gitlabnet/twofactorverify/client.go

@@ -26,7 +26,7 @@ type Response struct {
 type RequestBody struct {
 	KeyId      string `json:"key_id,omitempty"`
 	UserId     int64  `json:"user_id,omitempty"`
-	OTPAttempt string `json:"otp_attempt"`
+	OTPAttempt string `json:"otp_attempt,omitempty"`
 }
 
 func NewClient(config *config.Config) (*Client, error) {
@@ -44,7 +44,22 @@ func (c *Client) VerifyOTP(ctx context.Context, args *commandargs.Shell, otp str
 		return err
 	}
 
-	response, err := c.client.Post(ctx, "/two_factor_otp_check", requestBody)
+	response, err := c.client.Post(ctx, "/two_factor_manual_otp_check", requestBody)
+	if err != nil {
+		return err
+	}
+	defer response.Body.Close()
+
+	return parse(response)
+}
+
+func (c *Client) PushAuth(ctx context.Context, args *commandargs.Shell) error {
+	requestBody, err := c.getRequestBody(ctx, args, "")
+	if err != nil {
+		return err
+	}
+
+	response, err := c.client.Post(ctx, "/two_factor_push_otp_check", requestBody)
 	if err != nil {
 		return err
 	}

internal/gitlabnet/twofactorverify/client_test.go

@@ -17,49 +17,55 @@ import (
 )
 
 func initialize(t *testing.T) []testserver.TestRequestHandler {
+	handler := func(w http.ResponseWriter, r *http.Request) {
+		b, err := io.ReadAll(r.Body)
+		defer r.Body.Close()
+
+		require.NoError(t, err)
+
+		var requestBody *RequestBody
+		require.NoError(t, json.Unmarshal(b, &requestBody))
+
+		switch requestBody.KeyId {
+		case "0":
+			body := map[string]interface{}{
+				"success": true,
+			}
+			require.NoError(t, json.NewEncoder(w).Encode(body))
+		case "1":
+			body := map[string]interface{}{
+				"success": false,
+				"message": "error message",
+			}
+			require.NoError(t, json.NewEncoder(w).Encode(body))
+		case "2":
+			w.WriteHeader(http.StatusForbidden)
+			body := &client.ErrorResponse{
+				Message: "Not allowed!",
+			}
+			require.NoError(t, json.NewEncoder(w).Encode(body))
+		case "3":
+			w.Write([]byte("{ \"message\": \"broken json!\""))
+		case "4":
+			w.WriteHeader(http.StatusForbidden)
+		}
+
+		if requestBody.UserId == 1 {
+			body := map[string]interface{}{
+				"success": true,
+			}
+			require.NoError(t, json.NewEncoder(w).Encode(body))
+		}
+	}
+
 	requests := []testserver.TestRequestHandler{
 		{
-			Path: "/api/v4/internal/two_factor_otp_check",
-			Handler: func(w http.ResponseWriter, r *http.Request) {
-				b, err := io.ReadAll(r.Body)
-				defer r.Body.Close()
-
-				require.NoError(t, err)
-
-				var requestBody *RequestBody
-				require.NoError(t, json.Unmarshal(b, &requestBody))
-
-				switch requestBody.KeyId {
-				case "0":
-					body := map[string]interface{}{
-						"success": true,
-					}
-					require.NoError(t, json.NewEncoder(w).Encode(body))
-				case "1":
-					body := map[string]interface{}{
-						"success": false,
-						"message": "error message",
-					}
-					require.NoError(t, json.NewEncoder(w).Encode(body))
-				case "2":
-					w.WriteHeader(http.StatusForbidden)
-					body := &client.ErrorResponse{
-						Message: "Not allowed!",
-					}
-					require.NoError(t, json.NewEncoder(w).Encode(body))
-				case "3":
-					w.Write([]byte("{ \"message\": \"broken json!\""))
-				case "4":
-					w.WriteHeader(http.StatusForbidden)
-				}
-
-				if requestBody.UserId == 1 {
-					body := map[string]interface{}{
-						"success": true,
-					}
-					require.NoError(t, json.NewEncoder(w).Encode(body))
-				}
-			},
+			Path:    "/api/v4/internal/two_factor_manual_otp_check",
+			Handler: handler,
+		},
+		{
+			Path:    "/api/v4/internal/two_factor_push_otp_check",
+			Handler: handler,
 		},
 		{
 			Path: "/api/v4/internal/discover",
@@ -140,6 +146,57 @@ func TestErrorResponses(t *testing.T) {
 	}
 }
 
+func TestVerifyPush(t *testing.T) {
+	client := setup(t)
+
+	args := &commandargs.Shell{GitlabKeyId: "0"}
+	err := client.PushAuth(context.Background(), args)
+	require.NoError(t, err)
+}
+
+func TestErrorMessagePush(t *testing.T) {
+	client := setup(t)
+
+	args := &commandargs.Shell{GitlabKeyId: "1"}
+	err := client.PushAuth(context.Background(), args)
+	require.Equal(t, "error message", err.Error())
+}
+
+func TestErrorResponsesPush(t *testing.T) {
+	client := setup(t)
+
+	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) {
+			args := &commandargs.Shell{GitlabKeyId: tc.fakeId}
+			err := client.PushAuth(context.Background(), args)
+
+			require.EqualError(t, err, tc.expectedError)
+		})
+	}
+}
+
 func setup(t *testing.T) *Client {
 	requests := initialize(t)
 	url := testserver.StartSocketHttpServer(t, requests)

spec/gitlab_shell_two_factor_verify_spec.rb

@@ -11,71 +11,115 @@ describe 'bin/gitlab-shell 2fa_verify' do
       'SSH_ORIGINAL_COMMAND' => '2fa_verify' }
   end
 
+  let(:correct_otp) { '123456' }
+
   before(:context) do
     write_config('gitlab_url' => "http+unix://#{CGI.escape(tmp_socket_path)}")
   end
 
   def mock_server(server)
-    server.mount_proc('/api/v4/internal/two_factor_otp_check') do |req, res|
+    server.mount_proc('/api/v4/internal/two_factor_manual_otp_check') do |req, res|
       res.content_type = 'application/json'
       res.status = 200
 
       params = JSON.parse(req.body)
-      key_id = params['key_id'] || params['user_id'].to_s
 
-      if key_id == '100'
-        res.body = { success: true }.to_json
-      else
+      res.body = if params['otp_attempt'] == correct_otp
+                   { success: true }.to_json
+                 else
+                   { success: false, message: 'boom!' }.to_json
+                 end
+    end
+
+    server.mount_proc('/api/v4/internal/two_factor_push_otp_check') do |req, res|
+      res.content_type = 'application/json'
+      res.status = 200
+
+      params = JSON.parse(req.body)
+      id = params['key_id'] || params['user_id'].to_s
+
+      if id == '100'
         res.body = { success: false, message: 'boom!' }.to_json
+      else
+        res.body = { success: true }.to_json
       end
     end
 
-    server.mount_proc('/api/v4/internal/discover') do |_, res|
+    server.mount_proc('/api/v4/internal/discover') do |req, res|
       res.status = 200
       res.content_type = 'application/json'
-      res.body = { id: 100, name: 'Some User', username: 'someuser' }.to_json
+
+      if req.query['username'] == 'someone'
+        res.body = { id: 100, name: 'Some User', username: 'someuser' }.to_json
+      else
+        res.body = { id: 101, name: 'Another User', username: 'another' }.to_json
+      end
     end
   end
 
-  describe 'command' do
-    context 'when key is provided' do
-      let(:cmd) { "#{gitlab_shell_path} key-100" }
+  describe 'entering OTP manually' do
+    let(:cmd) { "#{gitlab_shell_path} key-100" }
 
-      it 'prints a successful verification message' do
-        verify_successful_verification!(cmd)
+    context 'when key is provided' do
+      it 'asks a user for a correct OTP' do
+        verify_successful_otp_verification!(cmd)
       end
     end
 
     context 'when username is provided' do
       let(:cmd) { "#{gitlab_shell_path} username-someone" }
 
-      it 'prints a successful verification message' do
-        verify_successful_verification!(cmd)
+      it 'asks a user for a correct OTP' do
+        verify_successful_otp_verification!(cmd)
       end
     end
 
-    context 'when API error occurs' do
+    it 'shows an error when an invalid otp is provided' do
+      Open3.popen2(env, cmd) do |stdin, stdout|
+        asks_for_otp(stdout)
+        stdin.puts('000000')
+
+        expect(stdout.flush.read).to eq("\nOTP validation failed: boom!\n")
+      end
+    end
+  end
+
+  describe 'authorizing via push' do
+    context 'when key is provided' do
       let(:cmd) { "#{gitlab_shell_path} key-101" }
 
-      it 'prints the error message' do
-        Open3.popen2(env, cmd) do |stdin, stdout|
-          expect(stdout.gets(5)).to eq('OTP: ')
+      it 'asks a user for a correct OTP' do
+        verify_successful_push_verification!(cmd)
+      end
+    end
 
-          stdin.puts('123456')
+    context 'when username is provided' do
+      let(:cmd) { "#{gitlab_shell_path} username-another" }
 
-          expect(stdout.flush.read).to eq("\nOTP validation failed.\nboom!\n")
-        end
+      it 'asks a user for a correct OTP' do
+        verify_successful_push_verification!(cmd)
       end
     end
   end
 
-  def verify_successful_verification!(cmd)
+  def verify_successful_otp_verification!(cmd)
     Open3.popen2(env, cmd) do |stdin, stdout|
-      expect(stdout.gets(5)).to eq('OTP: ')
-
-      stdin.puts('123456')
+      asks_for_otp(stdout)
+      stdin.puts(correct_otp)
 
       expect(stdout.flush.read).to eq("\nOTP validation successful. Git operations are now allowed.\n")
     end
   end
+
+  def verify_successful_push_verification!(cmd)
+    Open3.popen2(env, cmd) do |stdin, stdout|
+      asks_for_otp(stdout)
+
+      expect(stdout.flush.read).to eq("\nOTP has been validated by Push Authentication. Git operations are now allowed.\n")
+    end
+  end
+
+  def asks_for_otp(stdout)
+    expect(stdout.gets(5)).to eq('OTP: ')
+  end
 end