add callback for pipe start and pipe err (#145)

* add onstart onerr cb

* cover by e2e

* add missing plugin

* happy lint

cmd/sshpiperd/daemon.go

@@ -16,7 +16,7 @@ import (
 )
 
 type daemon struct {
-	config         *ssh.PiperConfig
+	config         *plugin.GrpcPluginConfig
 	lis            net.Listener
 	loginGraceTime time.Duration
 
@@ -25,7 +25,7 @@ type daemon struct {
 }
 
 func newDaemon(ctx *cli.Context) (*daemon, error) {
-	config := &ssh.PiperConfig{}
+	config := &plugin.GrpcPluginConfig{}
 	config.SetDefaults()
 
 	keybase64 := ctx.String("server-key-data")
@@ -139,7 +139,7 @@ func (d *daemon) run() error {
 			errorc := make(chan error)
 
 			go func() {
-				p, err := ssh.NewSSHPiperConn(c, d.config)
+				p, err := ssh.NewSSHPiperConn(c, &d.config.PiperConfig)
 
 				if err != nil {
 					errorc <- err
@@ -200,8 +200,16 @@ func (d *daemon) run() error {
 				}
 			}
 
+			if d.config.PipeStartCallback != nil {
+				d.config.PipeStartCallback(p.DownstreamConnMeta(), p.ChallengeContext())
+			}
+
 			err = p.WaitWithHook(uphook, downhook)
 
+			if d.config.PipeErrorCallback != nil {
+				d.config.PipeErrorCallback(p.DownstreamConnMeta(), p.ChallengeContext(), err)
+			}
+
 			log.Infof("connection from %v closed reason: %v", c.RemoteAddr(), err)
 		}(conn)
 	}

cmd/sshpiperd/internal/plugin/chain.go

@@ -10,7 +10,7 @@ import (
 )
 
 type ChainPlugins struct {
-	pluginsCallback []*ssh.PiperConfig
+	pluginsCallback []*GrpcPluginConfig
 	plugins         []*GrpcPlugin
 }
 
@@ -96,7 +96,7 @@ func (cp *ChainPlugins) NextAuthMethods(conn ssh.ConnMetadata, challengeCtx ssh.
 	return methods, nil
 }
 
-func (cp *ChainPlugins) InstallPiperConfig(config *ssh.PiperConfig) error {
+func (cp *ChainPlugins) InstallPiperConfig(config *GrpcPluginConfig) error {
 
 	config.CreateChallengeContext = func(conn ssh.ConnMetadata) (ssh.ChallengeContext, error) {
 		ctx, err := cp.CreateChallengeContext(conn)
@@ -140,5 +140,19 @@ func (cp *ChainPlugins) InstallPiperConfig(config *ssh.PiperConfig) error {
 		return ""
 	}
 
+	config.PipeStartCallback = func(conn ssh.ConnMetadata, challengeCtx ssh.ChallengeContext) {
+		cur := cp.pluginsCallback[challengeCtx.(*chainConnMeta).current]
+		if cur.PipeStartCallback != nil {
+			cur.PipeStartCallback(conn, challengeCtx)
+		}
+	}
+
+	config.PipeErrorCallback = func(conn ssh.ConnMetadata, challengeCtx ssh.ChallengeContext, err error) {
+		cur := cp.pluginsCallback[challengeCtx.(*chainConnMeta).current]
+		if cur.PipeErrorCallback != nil {
+			cur.PipeErrorCallback(conn, challengeCtx, err)
+		}
+	}
+
 	return nil
 }

cmd/sshpiperd/internal/plugin/grpc.go

@@ -19,6 +19,13 @@ import (
 	"google.golang.org/grpc/credentials/insecure"
 )
 
+type GrpcPluginConfig struct {
+	ssh.PiperConfig
+
+	PipeStartCallback func(conn ssh.ConnMetadata, challengeCtx ssh.ChallengeContext)
+	PipeErrorCallback func(conn ssh.ConnMetadata, challengeCtx ssh.ChallengeContext, err error)
+}
+
 type GrpcPlugin struct {
 	Name         string
 	OnNextPlugin func(conn ssh.ChallengeContext, upstream *libplugin.UpstreamNextPluginAuth) error
@@ -41,7 +48,7 @@ func DialGrpc(conn *grpc.ClientConn) (*GrpcPlugin, error) {
 	return p, nil
 }
 
-func (g *GrpcPlugin) InstallPiperConfig(config *ssh.PiperConfig) error {
+func (g *GrpcPlugin) InstallPiperConfig(config *GrpcPluginConfig) error {
 
 	cb, err := g.client.ListCallbacks(context.Background(), &libplugin.ListCallbackRequest{})
 	if err != nil {
@@ -119,6 +126,10 @@ func (g *GrpcPlugin) InstallPiperConfig(config *ssh.PiperConfig) error {
 			config.BannerCallback = g.BannerCallback
 		case "VerifyHostKey":
 			// ignore
+		case "PipeStart":
+			config.PipeStartCallback = g.PipeStartCallback
+		case "PipeError":
+			config.PipeErrorCallback = g.PipeErrorCallback
 		default:
 			return fmt.Errorf("unknown callback %s", c)
 		}
@@ -127,8 +138,8 @@ func (g *GrpcPlugin) InstallPiperConfig(config *ssh.PiperConfig) error {
 	return nil
 }
 
-func (g *GrpcPlugin) CreatePiperConfig() (*ssh.PiperConfig, error) {
-	config := &ssh.PiperConfig{}
+func (g *GrpcPlugin) CreatePiperConfig() (*GrpcPluginConfig, error) {
+	config := &GrpcPluginConfig{}
 	return config, g.InstallPiperConfig(config)
 }
 
@@ -474,6 +485,21 @@ func (g *GrpcPlugin) BannerCallback(conn ssh.ConnMetadata, challengeCtx ssh.Chal
 	return reply.GetMessage()
 }
 
+func (g *GrpcPlugin) PipeStartCallback(conn ssh.ConnMetadata, challengeCtx ssh.ChallengeContext) {
+	meta := toMeta(challengeCtx, conn)
+	_, _ = g.client.PipeStartNotice(context.Background(), &libplugin.PipeStartNoticeRequest{
+		Meta: meta,
+	})
+}
+
+func (g *GrpcPlugin) PipeErrorCallback(conn ssh.ConnMetadata, challengeCtx ssh.ChallengeContext, pipeerr error) {
+	meta := toMeta(challengeCtx, conn)
+	_, _ = g.client.PipeErrorNotice(context.Background(), &libplugin.PipeErrorNoticeRequest{
+		Meta:  meta,
+		Error: pipeerr.Error(),
+	})
+}
+
 func (g *GrpcPlugin) RecvLogs(writer io.Writer) error {
 	uid, err := uuid.NewRandom()
 	if err != nil {