From 1cd4e41a3ee0eb62a0d1f2ff43c6be40dd2129fe Mon Sep 17 00:00:00 2001 From: tgic Date: Mon, 29 Dec 2014 18:30:44 +0800 Subject: [PATCH] add user map api --- ssh/sshpiper.go | 13 +++++-- ssh/sshpiper_test.go | 85 ++++++++++++++++++++++++++++++++++++-------- 2 files changed, 81 insertions(+), 17 deletions(-) diff --git a/ssh/sshpiper.go b/ssh/sshpiper.go index dbd5a5a6..6bee6399 100644 --- a/ssh/sshpiper.go +++ b/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 diff --git a/ssh/sshpiper_test.go b/ssh/sshpiper_test.go index 583ee2ab..8c07c851 100644 --- a/ssh/sshpiper_test.go +++ b/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 }, })