chore: add auth cb to gswitch

pkg/gswitch/handlers.go

@@ -34,7 +34,16 @@ func (s *CLISession) handleLoginCommand() error {
 	if err != nil {
 		return err
 	}
-	if username != s.state.username || password != s.state.password {
+	if s.state.authCallback != nil {
+		err = s.state.authCallback(AuthRequest{
+			Method:   AuthMethodTelnet,
+			Username: username,
+			Password: password,
+		})
+		if err != nil {
+			return err
+		}
+	} else if username != s.state.username || password != s.state.password {
 		return fmt.Errorf("auth failed")
 	}
 	s.state.authenticated = true

pkg/gswitch/serve.go

@@ -11,12 +11,34 @@ import (
 	"golang.org/x/crypto/ssh"
 )
 
+type AuthMethod string
+
+const (
+	AuthMethodPassword  AuthMethod = "password"
+	AuthMethodPublicKey AuthMethod = "public_key"
+	AuthMethodTelnet    AuthMethod = "telnet"
+)
+
+// AuthRequest describes a gswitch authentication attempt.
+type AuthRequest struct {
+	Method    AuthMethod
+	Username  string
+	Password  string
+	PublicKey ssh.PublicKey
+}
+
+// AuthCallback checks a gswitch authentication attempt. Return nil to allow auth.
+type AuthCallback func(AuthRequest) error
+
 // SSHServerOptions configures the mock SSH switch listener.
 type SSHServerOptions struct {
 	Logger *zap.Logger
-	// Username and Password are used for password-based client auth.
+	// Username and Password are used for password-based client auth when AuthCallback is nil.
 	Username string
 	Password string
+	// AuthCallback, when set, is used to validate SSH password, SSH public-key,
+	// and telnet login/password authentication attempts.
+	AuthCallback AuthCallback
 	// ConnectionErrorProb is the probability (0–1) to drop a connection right after accept.
 	ConnectionErrorProb float64
 	// AuthorizedKeys, when non-empty, enables public-key client auth for Username.
@@ -35,15 +57,37 @@ func buildSSHServerConfig(opts SSHServerOptions) (*ssh.ServerConfig, error) {
 	config := &ssh.ServerConfig{
 		ServerVersion: "SSH-gswitch",
 		PasswordCallback: func(c ssh.ConnMetadata, pass []byte) (*ssh.Permissions, error) {
+			if opts.AuthCallback != nil {
+				err := opts.AuthCallback(AuthRequest{
+					Method:   AuthMethodPassword,
+					Username: c.User(),
+					Password: string(pass),
+				})
+				if err != nil {
+					return nil, err
+				}
+				return nil, nil
+			}
 			if c.User() == opts.Username && string(pass) == opts.Password {
 				return nil, nil
 			}
 			return nil, fmt.Errorf("password rejected for %q", c.User())
 		},
 	}
-	if len(opts.AuthorizedKeys) > 0 {
+	if opts.AuthCallback != nil || len(opts.AuthorizedKeys) > 0 {
 		allowed := opts.AuthorizedKeys
 		config.PublicKeyCallback = func(c ssh.ConnMetadata, pubKey ssh.PublicKey) (*ssh.Permissions, error) {
+			if opts.AuthCallback != nil {
+				err := opts.AuthCallback(AuthRequest{
+					Method:    AuthMethodPublicKey,
+					Username:  c.User(),
+					PublicKey: pubKey,
+				})
+				if err != nil {
+					return nil, err
+				}
+				return nil, nil
+			}
 			if c.User() != opts.Username {
 				return nil, fmt.Errorf("unknown user %q", c.User())
 			}
@@ -71,7 +115,7 @@ func ServeSSH(ctx context.Context, ln net.Listener, opts SSHServerOptions) error
 	if err != nil {
 		return err
 	}
-	conns := newConnections(opts.Username, opts.Password, opts.ConnectionErrorProb)
+	conns := newConnections(opts.Username, opts.Password, opts.AuthCallback, opts.ConnectionErrorProb)
 
 	for {
 		select {
@@ -110,7 +154,7 @@ func ServeSSH(ctx context.Context, ln net.Listener, opts SSHServerOptions) error
 // ServeTelnet accepts Telnet connections on ln until ln is closed or ctx is cancelled.
 func ServeTelnet(ctx context.Context, ln net.Listener, opts SSHServerOptions) error {
 	log := opts.logger()
-	conns := newConnections(opts.Username, opts.Password, opts.ConnectionErrorProb)
+	conns := newConnections(opts.Username, opts.Password, opts.AuthCallback, opts.ConnectionErrorProb)
 
 	for {
 		select {

pkg/gswitch/ssh_server.go

@@ -17,15 +17,17 @@ type connections struct {
 	done                atomic.Int32
 	username            string
 	password            string
+	authCallback        AuthCallback
 	connectionErrorProb float64
 }
 
-func newConnections(username, password string, connectionErrorProb float64) *connections {
+func newConnections(username, password string, authCallback AuthCallback, connectionErrorProb float64) *connections {
 	return &connections{
 		inFlight:            atomic.Int32{},
 		done:                atomic.Int32{},
 		username:            username,
 		password:            password,
+		authCallback:        authCallback,
 		connectionErrorProb: connectionErrorProb,
 	}
 }
@@ -136,7 +138,7 @@ func (c *connections) handleTelnetConnection(ctx context.Context, tcpConn net.Co
 	}
 
 	logger.Debug("start Telnet CLI session")
-	session := NewCLISessionWithAuth(telnetConn, c.username, c.password, logger)
+	session := newCLISessionWithAuth(telnetConn, c.username, c.password, c.authCallback, logger)
 	err = session.Run(ctx)
 	if err != nil {
 		logger.Debug("Telnet CLI session ended", zap.Error(err))

pkg/gswitch/states.go

@@ -34,6 +34,7 @@ type CLIState struct {
 	username      string
 	password      string
 	enablePass    string
+	authCallback  AuthCallback
 	authenticated bool
 	config        map[string]interface{}
 }
@@ -59,12 +60,17 @@ func NewCLIState(username, password string) *CLIState {
 }
 
 func NewCLIStateWithAuth(username, password string) *CLIState {
+	return newCLIStateWithAuth(username, password, nil)
+}
+
+func newCLIStateWithAuth(username, password string, authCallback AuthCallback) *CLIState {
 	return &CLIState{
 		mode:          ModeLogin,
 		hostname:      "switch",
 		username:      username,
 		password:      password,
 		enablePass:    password,
+		authCallback:  authCallback,
 		authenticated: false,
 		config:        make(map[string]interface{}),
 	}
@@ -92,9 +98,13 @@ func NewCLISession(conn ssh.Channel, username, password string, logger *zap.Logg
 }
 
 func NewCLISessionWithAuth(conn ssh.Channel, username, password string, logger *zap.Logger) *CLISession {
+	return newCLISessionWithAuth(conn, username, password, nil, logger)
+}
+
+func newCLISessionWithAuth(conn ssh.Channel, username, password string, authCallback AuthCallback, logger *zap.Logger) *CLISession {
 	return &CLISession{
 		conn:   conn,
-		state:  NewCLIStateWithAuth(username, password),
+		state:  newCLIStateWithAuth(username, password, authCallback),
 		logger: logger,
 		reader: bufio.NewReader(conn),
 		writer: conn,

pkg/server/server_integration_test.go

@@ -1,10 +1,12 @@
 package server_test
 
 import (
+	"bytes"
 	"context"
 	"crypto/ed25519"
 	"crypto/rand"
 	"encoding/pem"
+	"fmt"
 	"net"
 	"os"
 	"path/filepath"
@@ -52,11 +54,20 @@ func TestDevAuthPrivateKeyUsedForSSH(t *testing.T) {
 
 	go func() {
 		_ = gswitch.ServeSSH(ctx, ln, gswitch.SSHServerOptions{
-			Logger:              swLogger,
-			Username:            user,
-			Password:            "not-used-by-test",
+			Logger: swLogger,
+			AuthCallback: func(req gswitch.AuthRequest) error {
+				if req.Method != gswitch.AuthMethodPublicKey {
+					return fmt.Errorf("unexpected auth method %q", req.Method)
+				}
+				if req.Username != user {
+					return fmt.Errorf("unexpected user %q", req.Username)
+				}
+				if req.PublicKey == nil || !bytes.Equal(req.PublicKey.Marshal(), sshPub.Marshal()) {
+					return fmt.Errorf("unexpected public key")
+				}
+				return nil
+			},
 			ConnectionErrorProb: 0,
-			AuthorizedKeys:      []ssh.PublicKey{sshPub},
 		})
 	}()
 
@@ -99,9 +110,16 @@ func TestDevAuthLoginPasswordUsedForSSH(t *testing.T) {
 
 	go func() {
 		_ = gswitch.ServeSSH(ctx, ln, gswitch.SSHServerOptions{
-			Logger:              swLogger,
-			Username:            user,
-			Password:            pass,
+			Logger: swLogger,
+			AuthCallback: func(req gswitch.AuthRequest) error {
+				if req.Method != gswitch.AuthMethodPassword {
+					return fmt.Errorf("unexpected auth method %q", req.Method)
+				}
+				if req.Username != user || req.Password != pass {
+					return fmt.Errorf("bad credentials for %q", req.Username)
+				}
+				return nil
+			},
 			ConnectionErrorProb: 0,
 		})
 	}()
@@ -137,9 +155,16 @@ func TestDevAuthWrongPasswordSSHRejected(t *testing.T) {
 	const user = "switchadmin"
 	go func() {
 		_ = gswitch.ServeSSH(ctx, ln, gswitch.SSHServerOptions{
-			Logger:              zap.NewNop(),
-			Username:            user,
-			Password:            "server-real-pass",
+			Logger: zap.NewNop(),
+			AuthCallback: func(req gswitch.AuthRequest) error {
+				if req.Method != gswitch.AuthMethodPassword {
+					return fmt.Errorf("unexpected auth method %q", req.Method)
+				}
+				if req.Username != user || req.Password != "server-real-pass" {
+					return fmt.Errorf("bad credentials for %q", req.Username)
+				}
+				return nil
+			},
 			ConnectionErrorProb: 0,
 		})
 	}()