From 9d4b4cfe48e9121347993a4cd35be33c8004ff18 Mon Sep 17 00:00:00 2001
From: Andrew Lytvynov <andrew@goteleport.com>
Date: Wed, 19 May 2021 20:58:03 +0000
Subject: [PATCH] mfa: strip trailing newline when reading TOTP codes (#6948)

The newline should not be interpreted as part of the code (or any
`prompt.Input` result).
---
 lib/utils/prompt/confirmation.go      |  2 +-
 lib/utils/prompt/confirmation_test.go | 61 +++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 1 deletion(-)
 create mode 100644 lib/utils/prompt/confirmation_test.go

diff --git a/lib/utils/prompt/confirmation.go b/lib/utils/prompt/confirmation.go
index 3530cdeef7955..c594928d11e96 100644
--- a/lib/utils/prompt/confirmation.go
+++ b/lib/utils/prompt/confirmation.go
@@ -77,5 +77,5 @@ func Input(ctx context.Context, out io.Writer, in *ContextReader, question strin
 	if err != nil {
 		return "", trace.WrapWithMessage(err, "failed reading prompt response")
 	}
-	return string(answer), nil
+	return strings.TrimSpace(string(answer)), nil
 }
diff --git a/lib/utils/prompt/confirmation_test.go b/lib/utils/prompt/confirmation_test.go
new file mode 100644
index 0000000000000..20dcd3bf9153c
--- /dev/null
+++ b/lib/utils/prompt/confirmation_test.go
@@ -0,0 +1,61 @@
+/*
+Copyright 2021 Gravitational, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package prompt
+
+import (
+	"context"
+	"io"
+	"io/ioutil"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestInput(t *testing.T) {
+	t.Parallel()
+
+	out, in := io.Pipe()
+	t.Cleanup(func() { out.Close() })
+	write := func(t *testing.T, s string) {
+		_, err := in.Write([]byte(s))
+		require.NoError(t, err)
+	}
+
+	r := NewContextReader(out)
+	ctx := context.Background()
+
+	t.Run("no whitespace", func(t *testing.T) {
+		go write(t, "hi")
+		got, err := Input(ctx, ioutil.Discard, r, "")
+		require.NoError(t, err)
+		require.Equal(t, got, "hi")
+	})
+
+	t.Run("with whitespace", func(t *testing.T) {
+		go write(t, "hey\n")
+		got, err := Input(ctx, ioutil.Discard, r, "")
+		require.NoError(t, err)
+		require.Equal(t, got, "hey")
+	})
+
+	t.Run("closed input", func(t *testing.T) {
+		require.NoError(t, in.Close())
+		got, err := Input(ctx, ioutil.Discard, r, "")
+		require.ErrorIs(t, err, io.EOF)
+		require.Empty(t, got)
+	})
+}