-
-
Notifications
You must be signed in to change notification settings - Fork 150
cmd/sshpiperd: allow starting without plugins running #653
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: master
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -11,12 +11,14 @@ import ( | |
"os" | ||
"path" | ||
"path/filepath" | ||
"sync/atomic" | ||
"time" | ||
|
||
log "github.com/sirupsen/logrus" | ||
"github.com/tg123/sshpiper/cmd/sshpiperd/internal/plugin" | ||
"github.com/urfave/cli/v2" | ||
"golang.org/x/crypto/ssh" | ||
"tailscale.com/util/singleflight" | ||
) | ||
|
||
type daemon struct { | ||
|
@@ -29,8 +31,16 @@ type daemon struct { | |
usernameAsRecorddir bool | ||
filterHostkeysReqeust bool | ||
replyPing bool | ||
|
||
// Plugin initialization fields. | ||
pluginConfigs [][]string | ||
quit chan error | ||
pluginsInitialized atomic.Bool | ||
pluginIniSingleFlight singleflight.Group[unused, unused] | ||
} | ||
|
||
type unused struct{} | ||
|
||
func generateSshKey(keyfile string) error { | ||
_, privateKey, err := ed25519.GenerateKey(rand.Reader) | ||
if err != nil { | ||
|
@@ -222,12 +232,120 @@ func (d *daemon) install(plugins ...*plugin.GrpcPlugin) error { | |
return m.InstallPiperConfig(d.config) | ||
} | ||
|
||
func (d *daemon) setPluginConfigs(configs [][]string, quit chan error) { | ||
d.pluginConfigs = configs | ||
d.quit = quit | ||
} | ||
|
||
func (d *daemon) initializePlugins() error { | ||
var plugins []*plugin.GrpcPlugin | ||
|
||
for _, args := range d.pluginConfigs { | ||
var p *plugin.GrpcPlugin | ||
|
||
switch args[0] { | ||
case "grpc": | ||
log.Info("starting net grpc plugin: ") | ||
|
||
grpcplugin, err := createNetGrpcPlugin(args) | ||
if err != nil { | ||
return err | ||
} | ||
|
||
p = grpcplugin | ||
|
||
default: | ||
cmdplugin, err := createCmdPlugin(args) | ||
if err != nil { | ||
return err | ||
} | ||
|
||
go func() { | ||
d.quit <- <-cmdplugin.Quit | ||
}() | ||
|
||
p = &cmdplugin.GrpcPlugin | ||
} | ||
|
||
go func() { | ||
if err := p.RecvLogs(log.StandardLogger().Out); err != nil { | ||
log.Errorf("plugin %v recv logs error: %v", p.Name, err) | ||
} | ||
}() | ||
|
||
plugins = append(plugins, p) | ||
} | ||
|
||
if err := d.install(plugins...); err != nil { | ||
return err | ||
} | ||
return nil | ||
} | ||
|
||
// tryInitializePlugins attempts a single-flighted plugin initialization. | ||
func (d *daemon) tryInitializePlugins() error { | ||
if d.pluginsInitialized.Load() { | ||
return nil // previously initialized successfully | ||
} | ||
|
||
_, err, _ := d.pluginIniSingleFlight.Do(unused{}, func() (unused, error) { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. may i know if a simple lock works? |
||
err := d.initializePlugins() | ||
if err == nil { | ||
d.pluginsInitialized.Store(true) | ||
} | ||
return unused{}, err | ||
}) | ||
|
||
if err != nil { | ||
return fmt.Errorf("plugin initialization failed: %w", err) | ||
} | ||
return nil | ||
} | ||
|
||
// lazyPluginListener wraps a net.Listener to initialize plugins on first connection | ||
type lazyPluginListener struct { | ||
net.Listener | ||
daemon *daemon | ||
} | ||
|
||
func (l *lazyPluginListener) Accept() (net.Conn, error) { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. may i know the reason using a listener to active |
||
conn, err := l.Listener.Accept() | ||
if err != nil { | ||
return nil, err | ||
} | ||
if l.daemon.pluginsInitialized.Load() { | ||
return conn, nil | ||
} | ||
|
||
if err := l.daemon.tryInitializePlugins(); err != nil { | ||
log.Errorf("%s", err) | ||
conn.Close() | ||
return nil, err | ||
} | ||
|
||
return conn, nil | ||
} | ||
|
||
func (d *daemon) run() error { | ||
defer d.lis.Close() | ||
log.Infof("sshpiperd is listening on: %v", d.lis.Addr().String()) | ||
tcpAddr, ok := d.lis.Addr().(*net.TCPAddr) | ||
port := 0 | ||
if ok { | ||
port = tcpAddr.Port | ||
} | ||
log.WithFields(log.Fields{ | ||
"port": port, | ||
"addr": d.lis.Addr().String(), | ||
}).Info("sshpiperd is listening") | ||
|
||
// Wrap the listener with lazy plugin initialization | ||
lazyListener := &lazyPluginListener{ | ||
Listener: d.lis, | ||
daemon: d, | ||
} | ||
|
||
for { | ||
conn, err := d.lis.Accept() | ||
conn, err := lazyListener.Accept() | ||
if err != nil { | ||
log.Debugf("failed to accept connection: %v", err) | ||
continue | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
a bit worry about partial plugin failure and might create too many zombies
do not retry on inited plugin?