diff --git a/sshpiperd/upstream/database/pipemgr.go b/sshpiperd/upstream/database/pipemgr.go new file mode 100644 index 00000000..dd69887c --- /dev/null +++ b/sshpiperd/upstream/database/pipemgr.go @@ -0,0 +1,75 @@ +package database + +import ( + "fmt" + + "github.com/jinzhu/gorm" + + upstreamprovider "github.com/tg123/sshpiper/sshpiperd/upstream" +) + +func (p *plugin) ListPipe() ([]upstreamprovider.Pipe, error) { + db := p.db + + downstreams := make([]downstream, 0) + err := db.Set("gorm:auto_preload", true).Find(&downstreams).Error + if err != nil { + return nil, err + } + + pipes := make([]upstreamprovider.Pipe, 0, len(downstreams)) + for _, d := range downstreams { + + host, port, err := upstreamprovider.SplitHostPortForSSH(d.Upstream.Server.Address) + + if err != nil { + continue + } + + upuser := d.Upstream.Username + + if upuser == "" { + upuser = d.Username + } + + pipes = append(pipes, upstreamprovider.Pipe{ + Host: host, + Port: port, + Username: d.Username, + UpstreamUsername: upuser, + }) + } + + return pipes, nil +} + +func (p *plugin) CreatePipe(opt upstreamprovider.CreatePipeOption) error { + db := p.db + + return db.Create(&downstream{ + Username: opt.Username, + Upstream: upstream{ + Username: opt.UpstreamUsername, + AuthMapType: authMapTypeNone, + Server: server{ + Address: fmt.Sprintf("%v:%v", opt.Host, opt.Port), + IgnoreHostKey: true, + }, + }, + }).Error +} + +func (p *plugin) RemovePipe(name string) error { + db := p.db + + d, err := lookupDownstream(db, name) + if err != nil { + + if gorm.IsRecordNotFoundError(err) { + return nil + } + + return err + } + return db.Unscoped().Delete(d).Error +} diff --git a/sshpiperd/upstream/database/plugin.go b/sshpiperd/upstream/database/plugin.go index 94eb480e..b1055165 100644 --- a/sshpiperd/upstream/database/plugin.go +++ b/sshpiperd/upstream/database/plugin.go @@ -1,7 +1,6 @@ package database import ( - "fmt" "log" "github.com/jinzhu/gorm" @@ -21,72 +20,6 @@ type plugin struct { db *gorm.DB } -func (p *plugin) ListPipe() ([]upstreamprovider.Pipe, error) { - db := p.db - - downstreams := make([]downstream, 0) - err := db.Set("gorm:auto_preload", true).Find(&downstreams).Error - if err != nil { - return nil, err - } - - pipes := make([]upstreamprovider.Pipe, 0, len(downstreams)) - for _, d := range downstreams { - - host, port, err := upstreamprovider.SplitHostPortForSSH(d.Upstream.Server.Address) - - if err != nil { - continue - } - - upuser := d.Upstream.Username - - if upuser == "" { - upuser = d.Username - } - - pipes = append(pipes, upstreamprovider.Pipe{ - Host: host, - Port: port, - Username: d.Username, - UpstreamUsername: upuser, - }) - } - - return pipes, nil -} - -func (p *plugin) CreatePipe(opt upstreamprovider.CreatePipeOption) error { - db := p.db - - return db.Create(&downstream{ - Username: opt.Username, - Upstream: upstream{ - Username: opt.UpstreamUsername, - AuthMapType: authMapTypeNone, - Server: server{ - Address: fmt.Sprintf("%v:%v", opt.Host, opt.Port), - IgnoreHostKey: true, - }, - }, - }).Error -} - -func (p *plugin) RemovePipe(name string) error { - db := p.db - - d, err := lookupDownstream(db, name) - if err != nil { - - if gorm.IsRecordNotFoundError(err) { - return nil - } - - return err - } - return db.Unscoped().Delete(d).Error -} - func (p *plugin) GetHandler() upstreamprovider.Handler { return p.findUpstream } diff --git a/sshpiperd/upstream/provider.go b/sshpiperd/upstream/provider.go index 19a48bf0..947d41ac 100644 --- a/sshpiperd/upstream/provider.go +++ b/sshpiperd/upstream/provider.go @@ -30,12 +30,7 @@ type Pipe struct { Port int } -// Provider is a factory for Upstream Provider -type Provider interface { - registry.Plugin - - GetHandler() Handler - +type PipeManager interface { ListPipe() ([]Pipe, error) CreatePipe(opt CreatePipeOption) error @@ -43,6 +38,14 @@ type Provider interface { RemovePipe(name string) error } +// Provider is a factory for Upstream Provider +type Provider interface { + registry.Plugin + PipeManager + + GetHandler() Handler +} + var ( drivers = registry.NewRegistry() ) diff --git a/sshpiperd/upstream/workingdir/pipemgr.go b/sshpiperd/upstream/workingdir/pipemgr.go new file mode 100644 index 00000000..e3c3f795 --- /dev/null +++ b/sshpiperd/upstream/workingdir/pipemgr.go @@ -0,0 +1,72 @@ +package workingdir + +import ( + "fmt" + "io/ioutil" + "os" + + "github.com/tg123/sshpiper/sshpiperd/upstream" +) + +func (p *plugin) ListPipe() ([]upstream.Pipe, error) { + files, err := ioutil.ReadDir(config.WorkingDir) + if err != nil { + return nil, err + } + + pipes := make([]upstream.Pipe, 0, len(files)) + for _, file := range files { + if !file.IsDir() { + continue + } + + data, err := userUpstreamFile.read(file.Name()) + if err != nil { + continue + } + + host, port, mappedUser, err := parseUpstreamFile(string(data)) + + pipes = append(pipes, upstream.Pipe{ + Host: host, + Port: port, + Username: file.Name(), + UpstreamUsername: mappedUser, + }) + } + + return pipes, nil +} + +func (p *plugin) CreatePipe(opt upstream.CreatePipeOption) error { + err := os.MkdirAll(config.WorkingDir+"/"+opt.Username, 0775) + if err != nil { + return err + } + + path := userUpstreamFile.realPath(opt.Username) + if _, err := os.Stat(path); os.IsNotExist(err) { + + upuser := opt.UpstreamUsername + + if len(upuser) == 0 { + upuser = opt.Username + } + + content := fmt.Sprintf("%v@%v:%v", upuser, opt.Host, opt.Port) + return ioutil.WriteFile(path, []byte(content), 0600) + } else if err != nil { + return err + } + + return fmt.Errorf("upstream file of [%v] alreay exists", opt.Username) +} + +func (p *plugin) RemovePipe(name string) error { + path := userUpstreamFile.realPath(name) + if _, err := os.Stat(path); os.IsNotExist(err) { + return nil + } + + return os.Remove(path) +} diff --git a/sshpiperd/upstream/workingdir/plugin.go b/sshpiperd/upstream/workingdir/plugin.go index 7728b97b..48a35489 100644 --- a/sshpiperd/upstream/workingdir/plugin.go +++ b/sshpiperd/upstream/workingdir/plugin.go @@ -1,12 +1,8 @@ package workingdir import ( - "fmt" - "io/ioutil" - "log" - "os" - "github.com/tg123/sshpiper/sshpiperd/upstream" + "log" ) var logger *log.Logger @@ -14,69 +10,6 @@ var logger *log.Logger type plugin struct { } -func (p *plugin) ListPipe() ([]upstream.Pipe, error) { - files, err := ioutil.ReadDir(config.WorkingDir) - if err != nil { - return nil, err - } - - pipes := make([]upstream.Pipe, 0, len(files)) - for _, file := range files { - if !file.IsDir() { - continue - } - - data, err := userUpstreamFile.read(file.Name()) - if err != nil { - continue - } - - host, port, mappedUser, err := parseUpstreamFile(string(data)) - - pipes = append(pipes, upstream.Pipe{ - Host: host, - Port: port, - Username: file.Name(), - UpstreamUsername: mappedUser, - }) - } - - return pipes, nil -} - -func (p *plugin) CreatePipe(opt upstream.CreatePipeOption) error { - err := os.MkdirAll(config.WorkingDir+"/"+opt.Username, 0775) - if err != nil { - return err - } - - path := userUpstreamFile.realPath(opt.Username) - if _, err := os.Stat(path); os.IsNotExist(err) { - - upuser := opt.UpstreamUsername - - if len(upuser) == 0 { - upuser = opt.Username - } - - content := fmt.Sprintf("%v@%v:%v", upuser, opt.Host, opt.Port) - return ioutil.WriteFile(path, []byte(content), 0600) - } else if err != nil { - return err - } - - return fmt.Errorf("upstream file of [%v] alreay exists", opt.Username) -} - -func (p *plugin) RemovePipe(name string) error { - path := userUpstreamFile.realPath(name) - if _, err := os.Stat(path); os.IsNotExist(err) { - return nil - } - - return os.Remove(path) -} - func (p *plugin) GetName() string { return "workingdir" }