authmethods.patch

diff --git a/ssh/server.go b/ssh/server.go
index e73a1c1..9545076 100644
--- a/ssh/server.go
+++ b/ssh/server.go
@@ -62,6 +62,14 @@ type ServerConfig struct {
        // unknown.
        KeyboardInteractiveCallback func(conn ConnMetadata, client KeyboardInteractiveChallenge) (*Permissions, error)

+       // NextAuthMethodsCallback, if non-nil, is called each time the server
+       // replies with available authentication methods. Typically, this is
+       // called when an authentication request has failed.
+       //
+       // A PartialSuccess flag should be set to indicate
+       // a multi-step authentication process.
+       NextAuthMethodsCallback func(conn ConnMetadata) ([]string, bool, error)
+
        // AuthLogCallback, if non-nil, is called to log all authentication
        // attempts.
        AuthLogCallback func(conn ConnMetadata, method string, err error)
@@ -408,14 +416,23 @@ userAuthLoop:
                }

                var failureMsg userAuthFailureMsg
-               if config.PasswordCallback != nil {
-                       failureMsg.Methods = append(failureMsg.Methods, "password")
-               }
-               if config.PublicKeyCallback != nil {
-                       failureMsg.Methods = append(failureMsg.Methods, "publickey")
-               }
-               if config.KeyboardInteractiveCallback != nil {
-                       failureMsg.Methods = append(failureMsg.Methods, "keyboard-interactive")
+               if config.NextAuthMethodsCallback != nil {
+                       methods, partial, err := config.NextAuthMethodsCallback(s)
+                       if err != nil {
+                               return nil, err
+                       }
+                       failureMsg.Methods = methods
+                       failureMsg.PartialSuccess = partial
+               } else {
+                       if config.PasswordCallback != nil {
+                               failureMsg.Methods = append(failureMsg.Methods, "password")
+                       }
+                       if config.PublicKeyCallback != nil {
+                               failureMsg.Methods = append(failureMsg.Methods, "publickey")
+                       }
+                       if config.KeyboardInteractiveCallback != nil {
+                               failureMsg.Methods = append(failureMsg.Methods, "keyboard-interactive")
+                       }
                }

partialsuccess.patch

diff --git a/ssh/server.go b/ssh/server.go
index e73a1c1..1be9c70 100644
--- a/ssh/server.go
+++ b/ssh/server.go
@@ -259,6 +259,9 @@ func checkSourceAddress(addr net.Addr, sourceAddr string) error {
        return fmt.Errorf("ssh: remote address %v is not allowed because of source-address restriction", addr)
 }

+// PartialSuccess can be used if further authentication is required.
+var PartialSuccess = errors.New("authenticated with partial success")
+
 func (s *connection) serverAuthenticate(config *ServerConfig) (*Permissions, error) {
        var err error
        var cache pubKeyCache
@@ -422,6 +425,16 @@ userAuthLoop:
                        return nil, errors.New("ssh: no authentication methods configured but NoClientAuth is also false")
                }

+               if authErr == PartialSuccess {
+                       failureMsg.PartialSuccess = true
+                       for i := range failureMsg.Methods {
+                               if failureMsg.Methods[i] == userAuthReq.Method {
+                                       failureMsg.Methods = append(failureMsg.Methods[:i], failureMsg.Methods[i+1:]...)
+                                       break
+                               }
+                       }
+               }
+
                if err = s.transport.writePacket(Marshal(&failureMsg)); err != nil {
                        return nil, err
                }

The second one is simpler but the first gives more control to the user (customizing authentication order).