Enhance upstream banner handling with deduplication mode and refactor connection metadata (#597)

* Enhance upstream banner handling with deduplication mode and refactor connection metadata

* Add 'first-only' mode to upstream banner handling and update usage documentation
This commit is contained in:
Boshi Lian 2025-05-25 00:56:18 -07:00 committed by GitHub
parent 297b8c60e7
commit 2edfa6b4c3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 52 additions and 14 deletions

View file

@ -2,6 +2,7 @@ package main
import (
"crypto/ed25519"
"crypto/md5"
"crypto/rand"
"encoding/base64"
"encoding/pem"
@ -151,10 +152,46 @@ func newDaemon(ctx *cli.Context) (*daemon, error) {
case "passthrough":
// library will handle the banner to client
case "ignore":
config.UpstreamBannerCallback = func(downconn ssh.ServerPreAuthConn, banner string, challengeCtx ssh.ChallengeContext) error {
config.UpstreamBannerCallback = func(_ ssh.ServerPreAuthConn, _ string, _ ssh.ChallengeContext) error {
return nil
}
// case "dedup":
case "dedup":
config.UpstreamBannerCallback = func(downstream ssh.ServerPreAuthConn, banner string, ctx ssh.ChallengeContext) error {
meta, ok := ctx.Meta().(*plugin.PluginConnMeta)
if !ok {
// should not happen, but just in case
log.Warnf("upstream banner deduplication failed, cannot get plugin connection meta from challenge context")
return nil
}
hash := fmt.Sprintf("%x", md5.Sum([]byte(banner)))
key := fmt.Sprintf("sshpiperd.upstream.banner.%s", hash)
if meta.Metadata[key] == "true" {
return nil
}
meta.Metadata[key] = "true"
return downstream.SendAuthBanner(banner)
}
case "first-only":
config.UpstreamBannerCallback = func(downstream ssh.ServerPreAuthConn, banner string, ctx ssh.ChallengeContext) error {
meta, ok := ctx.Meta().(*plugin.PluginConnMeta)
if !ok {
// should not happen, but just in case
log.Warnf("upstream banner first-only failed, cannot get plugin connection meta from challenge context")
return nil
}
if meta.Metadata["sshpiperd.upstream.banner.sent"] == "true" {
return nil
}
meta.Metadata["sshpiperd.upstream.banner.sent"] = "true"
return downstream.SendAuthBanner(banner)
}
default:
return nil, fmt.Errorf("unknown upstream banner mode %q; allowed: 'passthrough' or 'ignore'", ctx.String("upstream-banner-mode"))
}

View file

@ -50,7 +50,7 @@ func (cp *ChainPlugins) onNextPlugin(challengeCtx ssh.ChallengeContext, upstream
}
type chainConnMeta struct {
connMeta
PluginConnMeta
current int
}
@ -61,15 +61,16 @@ func (cp *ChainPlugins) CreateChallengeContext(conn ssh.ServerPreAuthConn) (ssh.
}
meta := chainConnMeta{
connMeta: connMeta{
PluginConnMeta: PluginConnMeta{
UserName: conn.User(),
FromAddr: conn.RemoteAddr().String(),
UniqId: uiq.String(),
Metadata: make(map[string]string),
},
}
for _, p := range cp.plugins {
if err := p.NewConnection(&meta.connMeta); err != nil {
if err := p.NewConnection(&meta.PluginConnMeta); err != nil {
return nil, err
}
}

View file

@ -143,15 +143,15 @@ func (g *GrpcPlugin) CreatePiperConfig() (*GrpcPluginConfig, error) {
return config, g.InstallPiperConfig(config)
}
type connMeta libplugin.ConnMeta
type PluginConnMeta libplugin.ConnMeta
// ChallengedUsername implements ssh.ChallengeContext
func (m *connMeta) ChallengedUsername() string {
func (m *PluginConnMeta) ChallengedUsername() string {
return m.UserName
}
// Meta implements ssh.ChallengeContext
func (m *connMeta) Meta() interface{} {
func (m *PluginConnMeta) Meta() interface{} {
return m
}
@ -161,7 +161,7 @@ func (g *GrpcPlugin) CreateChallengeContext(conn ssh.ServerPreAuthConn) (ssh.Cha
return nil, err
}
meta := connMeta{
meta := PluginConnMeta{
UserName: conn.User(),
FromAddr: conn.RemoteAddr().String(),
UniqId: uiq.String(),
@ -171,7 +171,7 @@ func (g *GrpcPlugin) CreateChallengeContext(conn ssh.ServerPreAuthConn) (ssh.Cha
return &meta, g.NewConnection(&meta)
}
func (g *GrpcPlugin) NewConnection(meta *connMeta) error {
func (g *GrpcPlugin) NewConnection(meta *PluginConnMeta) error {
if g.hasNewConnectionCallback {
_, err := g.client.NewConnection(context.Background(), &libplugin.NewConnectionRequest{
Meta: &libplugin.ConnMeta{
@ -190,12 +190,12 @@ func (g *GrpcPlugin) NewConnection(meta *connMeta) error {
func toMeta(challengeCtx ssh.ChallengeContext, conn ssh.ConnMetadata) *libplugin.ConnMeta {
switch meta := challengeCtx.(type) {
case *connMeta:
case *PluginConnMeta:
meta.UserName = conn.User()
return (*libplugin.ConnMeta)(meta)
case *chainConnMeta:
meta.UserName = conn.User()
return (*libplugin.ConnMeta)(&meta.connMeta)
return (*libplugin.ConnMeta)(&meta.PluginConnMeta)
}
panic("unknown challenge context")
@ -596,7 +596,7 @@ func DialCmd(cmd *exec.Cmd) (*CmdPlugin, error) {
func GetUniqueID(ctx ssh.ChallengeContext) string {
switch meta := ctx.(type) {
case *connMeta:
case *PluginConnMeta:
return meta.UniqId
case *chainConnMeta:
return meta.UniqId

View file

@ -169,7 +169,7 @@ func main() {
&cli.StringFlag{
Name: "upstream-banner-mode",
Value: "passthrough",
Usage: "upstream banner mode, allowed values: 'passthrough' (pass the banner from upstream to client) or 'ignore' (ignore the banner from upstream)",
Usage: "upstream banner mode, allowed values: 'passthrough' (pass the banner from upstream to downstream), 'ignore' (ignore the banner from upstream), 'dedup' (deduplicate the banner from upstream, only pass same banner once to downstream), 'first-only' (only pass the first banner from upstream to downstream)",
EnvVars: []string{"SSHPIPERD_UPSTREAM_BANNER_MODE"},
},
&cli.BoolFlag{