From 2b43ac23a97504274e57944274b941f058ed5c52 Mon Sep 17 00:00:00 2001
From: Daniel Theophanes <kardianos@gmail.com>
Date: Mon, 30 Jan 2012 15:30:11 -0500
Subject: [PATCH] go.crypto/ssh: add User to ServerConn and pass *ServerConn to
 callbacks.

R=golang-dev, dave, agl
CC=golang-dev
https://golang.org/cl/5577070
---
 ssh/client_auth_test.go |  4 ++--
 ssh/server.go           | 15 +++++++++++----
 ssh/session_test.go     |  2 +-
 3 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/ssh/client_auth_test.go b/ssh/client_auth_test.go
index c41a93b..a55d8f4 100644
--- a/ssh/client_auth_test.go
+++ b/ssh/client_auth_test.go
@@ -119,10 +119,10 @@ var (
 	clientKeychain = new(keychain)
 	clientPassword = password("tiger")
 	serverConfig   = &ServerConfig{
-		PasswordCallback: func(user, pass string) bool {
+		PasswordCallback: func(conn *ServerConn, user, pass string) bool {
 			return user == "testuser" && pass == string(clientPassword)
 		},
-		PublicKeyCallback: func(user, algo string, pubkey []byte) bool {
+		PublicKeyCallback: func(conn *ServerConn, user, algo string, pubkey []byte) bool {
 			key := clientKeychain.keys[0].(*rsa.PrivateKey).PublicKey
 			expected := []byte(serializePublickey(key))
 			algoname := algoName(key)
diff --git a/ssh/server.go b/ssh/server.go
index 31011c6..1af0703 100644
--- a/ssh/server.go
+++ b/ssh/server.go
@@ -34,12 +34,12 @@ type ServerConfig struct {
 	// PasswordCallback, if non-nil, is called when a user attempts to
 	// authenticate using a password. It may be called concurrently from
 	// several goroutines.
-	PasswordCallback func(user, password string) bool
+	PasswordCallback func(conn *ServerConn, user, password string) bool
 
 	// PublicKeyCallback, if non-nil, is called when a client attempts public
 	// key authentication. It must return true iff the given public key is
 	// valid for the given user.
-	PublicKeyCallback func(user, algo string, pubkey []byte) bool
+	PublicKeyCallback func(conn *ServerConn, user, algo string, pubkey []byte) bool
 
 	// Cryptographic-related configuration.
 	Crypto CryptoConfig
@@ -150,6 +150,11 @@ type ServerConn struct {
 	// before attempting to authenticate with it, we end up with duplicate
 	// queries for public key validity.
 	cachedPubKeys []cachedPubKey
+
+	// User holds the successfully authenticated user name.
+	// It is empty if no authentication is used.  It is populated before
+	// any authentication callback is called and not assigned to after that.
+	User string
 }
 
 // Server returns a new SSH server connection
@@ -369,7 +374,7 @@ func (s *ServerConn) testPubKey(user, algo string, pubKey []byte) bool {
 		}
 	}
 
-	result := s.config.PublicKeyCallback(user, algo, pubKey)
+	result := s.config.PublicKeyCallback(s, user, algo, pubKey)
 	if len(s.cachedPubKeys) < maxCachedPubKeys {
 		c := cachedPubKey{
 			user:   user,
@@ -421,7 +426,8 @@ userAuthLoop:
 				return ParseError{msgUserAuthRequest}
 			}
 
-			if s.config.PasswordCallback(userAuthReq.User, string(password)) {
+			s.User = userAuthReq.User
+			if s.config.PasswordCallback(s, userAuthReq.User, string(password)) {
 				break userAuthLoop
 			}
 		case "publickey":
@@ -489,6 +495,7 @@ userAuthLoop:
 				default:
 					return errors.New("ssh: isAcceptableAlgo incorrect")
 				}
+				s.User = userAuthReq.User
 				if s.testPubKey(userAuthReq.User, algo, pubKey) {
 					break userAuthLoop
 				}
diff --git a/ssh/session_test.go b/ssh/session_test.go
index 4a3d22b..b314d14 100644
--- a/ssh/session_test.go
+++ b/ssh/session_test.go
@@ -18,7 +18,7 @@ type serverType func(*channel)
 // dial constructs a new test server and returns a *ClientConn.
 func dial(handler serverType, t *testing.T) *ClientConn {
 	pw := password("tiger")
-	serverConfig.PasswordCallback = func(user, pass string) bool {
+	serverConfig.PasswordCallback = func(conn *ServerConn, user, pass string) bool {
 		return user == "testuser" && pass == string(pw)
 	}
 	serverConfig.PublicKeyCallback = nil