add fallback to dababase handler

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 != "" {

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) {

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)"`
+}

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