Merge pull request #227 from gustavosbarreto/master

fix: bannerhandler

Author: gustavosbarreto

Diff for: server.go

@@ -39,6 +39,7 @@ type Server struct {
 	Version     string   // server version to be sent before the initial handshake
 	Banner      string   // server banner
 
+	BannerHandler                 BannerHandler                 // server banner handler, overrides Banner
 	KeyboardInteractiveHandler    KeyboardInteractiveHandler    // keyboard-interactive authentication handler
 	PasswordHandler               PasswordHandler               // password authentication handler
 	PublicKeyHandler              PublicKeyHandler              // public key authentication handler
@@ -134,10 +135,16 @@ func (srv *Server) config(ctx Context) *gossh.ServerConfig {
 		config.ServerVersion = "SSH-2.0-" + srv.Version
 	}
 	if srv.Banner != "" {
-		config.BannerCallback = func(conn gossh.ConnMetadata) string {
+		config.BannerCallback = func(_ gossh.ConnMetadata) string {
 			return srv.Banner
 		}
 	}
+	if srv.BannerHandler != nil {
+		config.BannerCallback = func(conn gossh.ConnMetadata) string {
+			applyConnMetadata(ctx, conn)
+			return srv.BannerHandler(ctx)
+		}
+	}
 	if srv.PasswordHandler != nil {
 		config.PasswordCallback = func(conn gossh.ConnMetadata, password []byte) (*gossh.Permissions, error) {
 			applyConnMetadata(ctx, conn)

Diff for: ssh.go

@@ -35,6 +35,9 @@ type Option func(*Server) error
 // Handler is a callback for handling established SSH sessions.
 type Handler func(Session)
 
+// BannerHandler is a callback for displaying the server banner.
+type BannerHandler func(ctx Context) string
+
 // PublicKeyHandler is a callback for performing public key authentication.
 type PublicKeyHandler func(ctx Context, key PublicKey) bool
 
@@ -115,8 +118,7 @@ func Handle(handler Handler) {
 
 // KeysEqual is constant time compare of the keys to avoid timing attacks.
 func KeysEqual(ak, bk PublicKey) bool {
-
-	//avoid panic if one of the keys is nil, return false instead
+	// avoid panic if one of the keys is nil, return false instead
 	if ak == nil || bk == nil {
 		return false
 	}