From 0c565bf13221fb55497d7ae2bb95694db1fd1bff Mon Sep 17 00:00:00 2001
From: Jamie Beverly <jamie.r.beverly@gmail.com>
Date: Sat, 25 Jun 2016 12:16:22 -0700
Subject: [PATCH] x/crypto/ssh: Add support for retryable authentication

Adds a new AuthMethod called "RetryableAuthMethod" which decorates any
other authmethod, allowing it to be retried up to maxTries before
aborting.

Fixes #16077

Change-Id: Ie310c24643e53dca4fa452750a69936674906484
Reviewed-on: https://go-review.googlesource.com/24156
Reviewed-by: Han-Wen Nienhuys <hanwen@google.com>
Run-TryBot: Han-Wen Nienhuys <hanwen@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
---
 ssh/client_auth.go      | 34 ++++++++++++++++++++++++++++++
 ssh/client_auth_test.go | 46 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+)

diff --git a/ssh/client_auth.go b/ssh/client_auth.go
index 6956ce4..294af0d 100644
--- a/ssh/client_auth.go
+++ b/ssh/client_auth.go
@@ -437,3 +437,37 @@ func (cb KeyboardInteractiveChallenge) auth(session []byte, user string, c packe
 		}
 	}
 }
+
+type retryableAuthMethod struct {
+	authMethod AuthMethod
+	maxTries   int
+}
+
+func (r *retryableAuthMethod) auth(session []byte, user string, c packetConn, rand io.Reader) (ok bool, methods []string, err error) {
+	for i := 0; r.maxTries <= 0 || i < r.maxTries; i++ {
+		ok, methods, err = r.authMethod.auth(session, user, c, rand)
+		if ok || err != nil { // either success or error terminate
+			return ok, methods, err
+		}
+	}
+	return ok, methods, err
+}
+
+func (r *retryableAuthMethod) method() string {
+	return r.authMethod.method()
+}
+
+// RetryableAuthMethod is a decorator for other auth methods enabling them to
+// be retried up to maxTries before considering that AuthMethod itself failed.
+// If maxTries is <= 0, will retry indefinitely
+//
+// This is useful for interactive clients using challenge/response type
+// authentication (e.g. Keyboard-Interactive, Password, etc) where the user
+// could mistype their response resulting in the server issuing a
+// SSH_MSG_USERAUTH_FAILURE (rfc4252 #8 [password] and rfc4256 #3.4
+// [keyboard-interactive]); Without this decorator, the non-retryable
+// AuthMethod would be removed from future consideration, and never tried again
+// (and so the user would never be able to retry their entry).
+func RetryableAuthMethod(auth AuthMethod, maxTries int) AuthMethod {
+	return &retryableAuthMethod{authMethod: auth, maxTries: maxTries}
+}
diff --git a/ssh/client_auth_test.go b/ssh/client_auth_test.go
index 2ea4462..2fa3103 100644
--- a/ssh/client_auth_test.go
+++ b/ssh/client_auth_test.go
@@ -391,3 +391,49 @@ func TestPermissionsPassing(t *testing.T) {
 func TestNoPermissionsPassing(t *testing.T) {
 	testPermissionsPassing(false, t)
 }
+
+func TestRetryableAuth(t *testing.T) {
+	n := 0
+	passwords := []string{"WRONG1", "WRONG2"}
+
+	config := &ClientConfig{
+		User: "testuser",
+		Auth: []AuthMethod{
+			RetryableAuthMethod(PasswordCallback(func() (string, error) {
+				p := passwords[n]
+				n++
+				return p, nil
+			}), 2),
+			PublicKeys(testSigners["rsa"]),
+		},
+	}
+
+	if err := tryAuth(t, config); err != nil {
+		t.Fatalf("unable to dial remote side: %s", err)
+	}
+	if n != 2 {
+		t.Fatalf("Did not try all passwords")
+	}
+}
+
+func ExampleRetryableAuthMethod(t *testing.T) {
+	user := "testuser"
+	NumberOfPrompts := 3
+
+	// Normally this would be a callback that prompts the user to answer the
+	// provided questions
+	Cb := func(user, instruction string, questions []string, echos []bool) (answers []string, err error) {
+		return []string{"answer1", "answer2"}, nil
+	}
+
+	config := &ClientConfig{
+		User: user,
+		Auth: []AuthMethod{
+			RetryableAuthMethod(KeyboardInteractiveChallenge(Cb), NumberOfPrompts),
+		},
+	}
+
+	if err := tryAuth(t, config); err != nil {
+		t.Fatalf("unable to dial remote side: %s", err)
+	}
+}