diff --git a/sshpiperd/upstream/database/handler.go b/sshpiperd/upstream/database/handler.go index 0e5ad39a..413bba64 100644 --- a/sshpiperd/upstream/database/handler.go +++ b/sshpiperd/upstream/database/handler.go @@ -2,6 +2,7 @@ package database import ( "bytes" + "github.com/jinzhu/gorm" "net" "golang.org/x/crypto/ssh" @@ -9,12 +10,9 @@ import ( func (p *plugin) findUpstream(conn ssh.ConnMetadata) (net.Conn, *ssh.AuthPipe, error) { - db := p.db - user := conn.User() + d, err := lookupDownstreamWithFallback(p.db, conn.User()) - d := downstream{} - - if err := db.Set("gorm:auto_preload", true).Where(&downstream{Username: user}).First(&d).Error; err != nil { + if err != nil { return nil, nil, err } @@ -60,6 +58,40 @@ func (p *plugin) findUpstream(conn ssh.ConnMetadata) (net.Conn, *ssh.AuthPipe, e return c, &pipe, nil } +func lookupDownstreamWithFallback(db *gorm.DB, user string) (*downstream, error) { + d, err := lookupDownstream(db, user) + + if gorm.IsRecordNotFoundError(err) { + fallback, _ := lookupConfigValue(db, fallbackUserEntry) + + if len(fallback) > 0 { + return lookupDownstream(db, fallback) + } + } + + return d, err +} + +func lookupDownstream(db *gorm.DB, user string) (*downstream, error) { + d := downstream{} + + if err := db.Set("gorm:auto_preload", true).Where(&downstream{Username: user}).First(&d).Error; err != nil { + + return nil, err + } + + return &d, nil +} + +func lookupConfigValue(db *gorm.DB, entry string) (string, error) { + c := config{} + if err := db.Where(&config{Entry: entry}).First(&c).Error; err != nil { + return "", err + } + + return c.Value, nil +} + func dial(addr string) (net.Conn, error) { if _, _, err := net.SplitHostPort(addr); err != nil && addr != "" { diff --git a/sshpiperd/upstream/database/handler_test.go b/sshpiperd/upstream/database/handler_test.go index 8c1027b2..5be7b01c 100644 --- a/sshpiperd/upstream/database/handler_test.go +++ b/sshpiperd/upstream/database/handler_test.go @@ -139,6 +139,21 @@ func TestFindUpstream(t *testing.T) { if auth.User != "findup0" { t.Error("auth pipe user name is not correct") } + + db.Delete(&config{Entry: fallbackUserEntry}) + _, _, err = h(testconn{"not_exists"}) + if err == nil { + t.Error("should not found any user") + } + + db.NewRecord(&config{Entry: fallbackUserEntry, Value: "findup0"}) + _, auth, err = h(testconn{"finddown0"}) + if err != nil { + t.Fatal(err) + } + if auth.User != "findup0" { + t.Error("auth pipe user name is not correct (fallback)") + } } func TestPublicKeyCallback(t *testing.T) { diff --git a/sshpiperd/upstream/database/model.go b/sshpiperd/upstream/database/model.go index 869acd4e..8da8beda 100644 --- a/sshpiperd/upstream/database/model.go +++ b/sshpiperd/upstream/database/model.go @@ -12,6 +12,8 @@ const ( authMapTypePrivateKey ) +const fallbackUserEntry = "FALLBACK_USER" + type keydata struct { gorm.Model @@ -77,3 +79,10 @@ type downstream struct { AuthorizedKeys []authorizedKey } + +type config struct { + gorm.Model + + Entry string `gorm:"type:varchar(45);unique_index"` + Value string `gorm:"type:varchar(100)"` +} diff --git a/sshpiperd/upstream/database/plugin.go b/sshpiperd/upstream/database/plugin.go index 44979449..186b73da 100644 --- a/sshpiperd/upstream/database/plugin.go +++ b/sshpiperd/upstream/database/plugin.go @@ -33,7 +33,6 @@ func (p *plugin) Init(glogger *log.Logger) error { if err != nil { return err } - //db.LogMode(true) db.AutoMigrate( new(keydata), @@ -43,6 +42,7 @@ func (p *plugin) Init(glogger *log.Logger) error { new(upstream), new(authorizedKey), new(downstream), + new(config), ) p.db = db