add user map api

ssh/sshpiper.go

@@ -25,9 +25,10 @@ type SSHPiperConfig struct {
 
 	// FindUpstream, must not be nil, is called when SSHPiper decided to establish a
 	// ssh connection to upstream server.  a connection, net.Conn, to upstream
-	// should be returned.
+	// and upstream username should be returned.
+	// SSHPiper will use the username from downstream if empty username is returned.
 	// If any error occurs, the piped connection will be closed.
-	FindUpstream func(conn ConnMetadata) (net.Conn, error)
+	FindUpstream func(conn ConnMetadata) (net.Conn, string, error)
 
 	// MapPublicKey, if non-nil, is called when downstream requests a publickey auth.
 	// SSHPiper will sign the auth packet message using the returned Signer.
@@ -147,13 +148,17 @@ func NewSSHPiperConn(conn net.Conn, piper *SSHPiperConfig) (pipe *SSHPiperConn,
 		}
 	}
 
-	upconn, err := piper.FindUpstream(d)
+	upconn, mappedUser, err := piper.FindUpstream(d)
 	if err != nil {
 		return nil, err
 	}
 
 	addr := upconn.RemoteAddr().String()
 
+	if mappedUser == "" {
+		mappedUser = d.user
+	}
+
 	u, err := newUpstream(upconn, addr, &ClientConfig{})
 	if err != nil {
 		return nil, err
@@ -171,6 +176,8 @@ func NewSSHPiperConn(conn net.Conn, piper *SSHPiperConfig) (pipe *SSHPiperConn,
 
 	p.processAuthMsg = func(msg *userAuthRequestMsg) (*userAuthRequestMsg, error) {
 
+		msg.User = mappedUser
+
 		// only public msg need
 		if msg.Method != "publickey" || piper.MapPublicKey == nil {
 			return msg, nil

ssh/sshpiper_test.go

@@ -22,13 +22,14 @@ func ExampleNewSSHPiperConn() {
 
 	piper := &SSHPiperConfig{
 		// return conn dial to serverAddr
-		FindUpstream: func(conn ConnMetadata) (net.Conn, error) {
+		FindUpstream: func(conn ConnMetadata) (net.Conn, string, error) {
 			c, err := net.Dial("tcp", serverAddr)
 			if err != nil {
-				return nil, err
+				return nil, "", err
 			}
 
-			return c, nil
+			// change upstream username to root
+			return c, "root", nil
 		},
 	}
 
@@ -98,14 +99,24 @@ func TestFindUpstreamCallback(t *testing.T) {
 	var called bool
 
 	c, err := dialPiper(&SSHPiperConfig{
-		FindUpstream: func(conn ConnMetadata) (net.Conn, error) {
-
-			called = true
+		FindUpstream: func(conn ConnMetadata) (net.Conn, string, error) {
 			if username != conn.User() {
 				t.Errorf("different username")
 			}
 
-			return nil, fmt.Errorf("not impl")
+			s, err := dialUpstream(simpleEchoHandler, &ServerConfig{
+				PasswordCallback: func(conn ConnMetadata, password []byte) (*Permissions, error) {
+					called = true
+
+					if conn.User() != username {
+						t.Errorf("default username changed")
+					}
+
+					return nil, nil
+				},
+			}, t)
+
+			return s, "", err
 		},
 	})
 
@@ -113,7 +124,10 @@ func TestFindUpstreamCallback(t *testing.T) {
 		t.Fatalf("connect dial to piper: %v", err)
 	}
 
-	NewClientConn(c, "", &ClientConfig{User: username})
+	NewClientConn(c, "", &ClientConfig{
+		User: username,
+		Auth: []AuthMethod{Password("password")},
+	})
 
 	if !called {
 		t.Fatalf("FindUpstream not called")
@@ -121,6 +135,46 @@ func TestFindUpstreamCallback(t *testing.T) {
 	}
 }
 
+// TODO clean up duplicate code
+func TestFindUpstreamWithUserCallback(t *testing.T) {
+	const username = "testuser"
+	const mappedname = "mappedname"
+
+	var called bool
+
+	c, err := dialPiper(&SSHPiperConfig{
+		FindUpstream: func(conn ConnMetadata) (net.Conn, string, error) {
+
+			s, err := dialUpstream(simpleEchoHandler, &ServerConfig{
+				PasswordCallback: func(conn ConnMetadata, password []byte) (*Permissions, error) {
+					called = true
+
+					if conn.User() != mappedname {
+						t.Errorf("bad mapped username")
+					}
+
+					return nil, nil
+				},
+			}, t)
+
+			return s, mappedname, err
+		},
+	})
+
+	if err != nil {
+		t.Fatalf("connect dial to piper: %v", err)
+	}
+
+	NewClientConn(c, "", &ClientConfig{
+		User: username,
+		Auth: []AuthMethod{Password("password")},
+	})
+
+	if !called {
+		t.Fatalf("FindUpstream not called")
+	}
+}
+
 func TestMapPublicKey(t *testing.T) {
 
 	certChecker := CertChecker{
@@ -140,10 +194,11 @@ func TestMapPublicKey(t *testing.T) {
 	}
 
 	c, err := dialPiper(&SSHPiperConfig{
-		FindUpstream: func(conn ConnMetadata) (net.Conn, error) {
-			return dialUpstream(simpleEchoHandler, &ServerConfig{
+		FindUpstream: func(conn ConnMetadata) (net.Conn, string, error) {
+			s, err := dialUpstream(simpleEchoHandler, &ServerConfig{
 				PublicKeyCallback: certChecker.Authenticate,
 			}, t)
+			return s, "", err
 		},
 
 		MapPublicKey: func(conn ConnMetadata, key PublicKey) (Signer, error) {
@@ -186,8 +241,9 @@ func TestAdditionalChallenge(t *testing.T) {
 			}
 			return false, fmt.Errorf("keyboard-interactive failed")
 		},
-		FindUpstream: func(conn ConnMetadata) (net.Conn, error) {
-			return dialUpstream(simpleEchoHandler, &ServerConfig{NoClientAuth: true}, t)
+		FindUpstream: func(conn ConnMetadata) (net.Conn, string, error) {
+			s, err := dialUpstream(simpleEchoHandler, &ServerConfig{NoClientAuth: true}, t)
+			return s, "", err
 		},
 	})
 
@@ -257,8 +313,9 @@ func dialUpstream(handler serverType, upstream *ServerConfig, t *testing.T) (net
 func TestPipeData(t *testing.T) {
 
 	c, err := dialPiper(&SSHPiperConfig{
-		FindUpstream: func(conn ConnMetadata) (net.Conn, error) {
-			return dialUpstream(simpleEchoHandler, &ServerConfig{NoClientAuth: true}, t)
+		FindUpstream: func(conn ConnMetadata) (net.Conn, string, error) {
+			s, err := dialUpstream(simpleEchoHandler, &ServerConfig{NoClientAuth: true}, t)
+			return s, "", err
 		},
 	})