Compare commits

..

No commits in common. "master" and "v0.1.1" have entirely different histories.

6 changed files with 84 additions and 568 deletions

View File

@ -4,9 +4,7 @@ import (
"errors"
"fmt"
"io"
"net"
"strings"
"sync"
"time"
"golang.org/x/crypto/ssh"
@ -17,53 +15,24 @@ var requestSSHAgentForwarding = func(session *ssh.Session) error {
return sshagent.RequestAgentForwarding(session)
}
const sshAgentChannelType = "auth-agent@openssh.com"
var routeSSHAgentForwarding = func(client *ssh.Client, timeout time.Duration) (io.Closer, error) {
return startSSHAgentForwardProxy(client, timeout)
var routeSSHAgentForwarding = func(client *ssh.Client, keyring sshagent.Agent) error {
return sshagent.ForwardToAgent(client, keyring)
}
var probeSSHAgentForwarding = func(timeout time.Duration) error {
var newSSHAgentForwarder = func(timeout time.Duration) (sshagent.Agent, io.Closer, error) {
conn, err := dialSSHAgent(timeout)
if err != nil {
return wrapSSHAgentForwardingUnavailable(err)
return nil, nil, wrapSSHAgentForwardingUnavailable(err)
}
if conn == nil {
return wrapSSHAgentForwardingUnavailable(errors.New("empty agent connection"))
return nil, nil, wrapSSHAgentForwardingUnavailable(errors.New("empty agent connection"))
}
return conn.Close()
return sshagent.NewClient(conn), conn, nil
}
var errSSHAgentForwardingDenied = errors.New("ssh-agent forwarding request denied")
var errSSHAgentForwardingUnavailable = errors.New("ssh-agent forwarding unavailable")
type sshAgentForwardProxy struct {
stopOnce sync.Once
stopCh chan struct{}
activeMu sync.Mutex
active map[*sshAgentForwardBridge]struct{}
}
func (p *sshAgentForwardProxy) Close() error {
if p == nil {
return nil
}
p.stopOnce.Do(func() {
close(p.stopCh)
})
p.closeActive()
return nil
}
type sshAgentForwardBridge struct {
proxy *sshAgentForwardProxy
channel ssh.Channel
conn net.Conn
closeOnce sync.Once
}
func (s *StarSSH) RequestAgentForwarding(session *ssh.Session) error {
if s == nil {
return errors.New("ssh client is nil")
@ -111,21 +80,20 @@ func (s *StarSSH) ensureAgentForwarding() error {
return err
}
timeout := effectiveDialTimeout(s.LoginInfo)
if err := probeSSHAgentForwarding(timeout); err != nil {
keyring, closer, err := newSSHAgentForwarder(s.LoginInfo.Timeout)
if err != nil {
return wrapSSHAgentForwardingUnavailable(err)
}
if s.closing.Load() {
_ = closer.Close()
return errSSHClientClosing
}
closer, err := routeSSHAgentForwarding(client, timeout)
if err != nil {
if err := routeSSHAgentForwarding(client, keyring); err != nil {
_ = closer.Close()
return err
}
if !s.canAttachAgentForwarder(client) {
if closer != nil {
_ = closer.Close()
}
_ = closer.Close()
return errSSHClientClosing
}
s.agentForwarder = closer
@ -181,175 +149,3 @@ func wrapSSHAgentForwardingUnavailable(err error) error {
}
return fmt.Errorf("%w: %v", errSSHAgentForwardingUnavailable, err)
}
func startSSHAgentForwardProxy(client *ssh.Client, timeout time.Duration) (io.Closer, error) {
if client == nil {
return nil, errors.New("ssh client is nil")
}
channels := client.HandleChannelOpen(sshAgentChannelType)
if channels == nil {
return nil, errors.New("agent: already have handler for " + sshAgentChannelType)
}
proxy := &sshAgentForwardProxy{
stopCh: make(chan struct{}),
active: make(map[*sshAgentForwardBridge]struct{}),
}
go func() {
for {
select {
case <-proxy.stopCh:
return
case ch, ok := <-channels:
if !ok {
return
}
go handleSSHAgentForwardChannel(proxy, ch, timeout)
}
}
}()
return proxy, nil
}
func handleSSHAgentForwardChannel(proxy *sshAgentForwardProxy, ch ssh.NewChannel, timeout time.Duration) {
if ch == nil {
return
}
conn, err := dialSSHAgent(timeout)
if err != nil {
_ = ch.Reject(ssh.ConnectionFailed, err.Error())
return
}
if conn == nil {
_ = ch.Reject(ssh.ConnectionFailed, "ssh-agent connection unavailable")
return
}
channel, reqs, err := ch.Accept()
if err != nil {
_ = conn.Close()
return
}
go ssh.DiscardRequests(reqs)
bridge := &sshAgentForwardBridge{
proxy: proxy,
channel: channel,
conn: conn,
}
if !proxy.registerBridge(bridge) {
bridge.close()
return
}
go bridge.run()
}
func proxySSHAgentChannel(channel ssh.Channel, conn net.Conn) {
bridge := &sshAgentForwardBridge{
channel: channel,
conn: conn,
}
bridge.run()
}
func (b *sshAgentForwardBridge) run() {
if b == nil {
return
}
defer b.unregister()
var wg sync.WaitGroup
wg.Add(2)
go func() {
defer wg.Done()
_, _ = io.Copy(b.channel, b.conn)
b.close()
}()
go func() {
defer wg.Done()
_, _ = io.Copy(b.conn, b.channel)
b.close()
}()
wg.Wait()
}
func (b *sshAgentForwardBridge) close() {
if b == nil {
return
}
b.closeOnce.Do(func() {
closeWriter(b.channel)
closeWriter(b.conn)
if b.channel != nil {
_ = b.channel.Close()
}
if b.conn != nil {
_ = b.conn.Close()
}
})
}
func (b *sshAgentForwardBridge) unregister() {
if b == nil || b.proxy == nil {
return
}
b.proxy.unregisterBridge(b)
}
func (p *sshAgentForwardProxy) registerBridge(bridge *sshAgentForwardBridge) bool {
if p == nil || bridge == nil {
return false
}
p.activeMu.Lock()
defer p.activeMu.Unlock()
select {
case <-p.stopCh:
return false
default:
}
if p.active == nil {
p.active = make(map[*sshAgentForwardBridge]struct{})
}
p.active[bridge] = struct{}{}
return true
}
func (p *sshAgentForwardProxy) unregisterBridge(bridge *sshAgentForwardBridge) {
if p == nil || bridge == nil {
return
}
p.activeMu.Lock()
defer p.activeMu.Unlock()
delete(p.active, bridge)
}
func (p *sshAgentForwardProxy) closeActive() {
if p == nil {
return
}
p.activeMu.Lock()
active := make([]*sshAgentForwardBridge, 0, len(p.active))
for bridge := range p.active {
active = append(active, bridge)
}
p.active = make(map[*sshAgentForwardBridge]struct{})
p.activeMu.Unlock()
for _, bridge := range active {
bridge.close()
}
}
func closeWriter(value any) {
type closeWriter interface {
CloseWrite() error
}
if cw, ok := value.(closeWriter); ok {
_ = cw.CloseWrite()
}
}

View File

@ -1,16 +1,14 @@
package starssh
import (
"bytes"
"errors"
"io"
"net"
"sync"
"sync/atomic"
"testing"
"time"
"golang.org/x/crypto/ssh"
sshagent "golang.org/x/crypto/ssh/agent"
)
type testCloser struct {
@ -22,90 +20,15 @@ func (c *testCloser) Close() error {
return nil
}
type trackedConn struct {
net.Conn
closed atomic.Int32
}
func (c *trackedConn) Close() error {
c.closed.Add(1)
if c.Conn == nil {
return nil
}
return c.Conn.Close()
}
type testSSHChannel struct {
readFunc func([]byte) (int, error)
stderr bytes.Buffer
closed atomic.Int32
closeOnce sync.Once
closeCh chan struct{}
}
func newTestSSHChannel(readFunc func([]byte) (int, error)) *testSSHChannel {
return &testSSHChannel{
readFunc: readFunc,
closeCh: make(chan struct{}),
}
}
func newBlockingTestSSHChannel() *testSSHChannel {
ch := newTestSSHChannel(nil)
ch.readFunc = func(p []byte) (int, error) {
<-ch.closeCh
return 0, io.EOF
}
return ch
}
func (c *testSSHChannel) Read(p []byte) (int, error) {
if c == nil {
return 0, io.EOF
}
if c.readFunc != nil {
return c.readFunc(p)
}
return 0, io.EOF
}
func (c *testSSHChannel) Write(p []byte) (int, error) {
return len(p), nil
}
func (c *testSSHChannel) Close() error {
if c == nil {
return nil
}
c.closeOnce.Do(func() {
c.closed.Add(1)
close(c.closeCh)
})
return nil
}
func (c *testSSHChannel) CloseWrite() error {
return nil
}
func (c *testSSHChannel) SendRequest(name string, wantReply bool, payload []byte) (bool, error) {
return false, nil
}
func (c *testSSHChannel) Stderr() io.ReadWriter {
return &c.stderr
}
func TestNewExecSessionEnablesAgentForwardingOnce(t *testing.T) {
oldNewSSHSession := newSSHSession
oldProbeSSHAgentForwarding := probeSSHAgentForwarding
oldNewSSHAgentForwarder := newSSHAgentForwarder
oldRouteSSHAgentForwarding := routeSSHAgentForwarding
oldRequestSSHAgentForwarding := requestSSHAgentForwarding
oldCloseSSHClient := closeSSHClient
t.Cleanup(func() {
newSSHSession = oldNewSSHSession
probeSSHAgentForwarding = oldProbeSSHAgentForwarding
newSSHAgentForwarder = oldNewSSHAgentForwarder
routeSSHAgentForwarding = oldRouteSSHAgentForwarding
requestSSHAgentForwarding = oldRequestSSHAgentForwarding
closeSSHClient = oldCloseSSHClient
@ -127,26 +50,26 @@ func TestNewExecSessionEnablesAgentForwardingOnce(t *testing.T) {
return &ssh.Session{}, nil
}
var probeCalls atomic.Int32
var agentInitCalls atomic.Int32
closer := &testCloser{}
probeSSHAgentForwarding = func(timeout time.Duration) error {
probeCalls.Add(1)
newSSHAgentForwarder = func(timeout time.Duration) (sshagent.Agent, io.Closer, error) {
agentInitCalls.Add(1)
if timeout != time.Second {
t.Fatalf("unexpected forwarding timeout: %v", timeout)
}
return nil
return sshagent.NewKeyring(), closer, nil
}
var routeCalls atomic.Int32
routeSSHAgentForwarding = func(client *ssh.Client, timeout time.Duration) (io.Closer, error) {
routeSSHAgentForwarding = func(client *ssh.Client, keyring sshagent.Agent) error {
routeCalls.Add(1)
if client != baseClient {
t.Fatalf("unexpected routed client %p", client)
}
if timeout != time.Second {
t.Fatalf("unexpected routed timeout: %v", timeout)
if keyring == nil {
t.Fatal("expected non-nil forwarded agent keyring")
}
return closer, nil
return nil
}
var requestCalls atomic.Int32
@ -165,8 +88,8 @@ func TestNewExecSessionEnablesAgentForwardingOnce(t *testing.T) {
t.Fatalf("second exec session: %v", err)
}
if probeCalls.Load() != 1 {
t.Fatalf("expected one agent probe, got %d", probeCalls.Load())
if agentInitCalls.Load() != 1 {
t.Fatalf("expected one agent forwarder init, got %d", agentInitCalls.Load())
}
if routeCalls.Load() != 1 {
t.Fatalf("expected one agent route registration, got %d", routeCalls.Load())
@ -187,13 +110,13 @@ func TestNewExecSessionEnablesAgentForwardingOnce(t *testing.T) {
func TestNewPTYSessionEnablesAgentForwardingWhenConfigured(t *testing.T) {
oldNewSSHSession := newSSHSession
oldRequestSessionPTY := requestSessionPTY
oldProbeSSHAgentForwarding := probeSSHAgentForwarding
oldNewSSHAgentForwarder := newSSHAgentForwarder
oldRouteSSHAgentForwarding := routeSSHAgentForwarding
oldRequestSSHAgentForwarding := requestSSHAgentForwarding
t.Cleanup(func() {
newSSHSession = oldNewSSHSession
requestSessionPTY = oldRequestSessionPTY
probeSSHAgentForwarding = oldProbeSSHAgentForwarding
newSSHAgentForwarder = oldNewSSHAgentForwarder
routeSSHAgentForwarding = oldRouteSSHAgentForwarding
requestSSHAgentForwarding = oldRequestSSHAgentForwarding
})
@ -215,12 +138,10 @@ func TestNewPTYSessionEnablesAgentForwardingWhenConfigured(t *testing.T) {
return nil
}
probeSSHAgentForwarding = func(timeout time.Duration) error {
return nil
}
routeSSHAgentForwarding = func(client *ssh.Client, timeout time.Duration) (io.Closer, error) {
return &testCloser{}, nil
newSSHAgentForwarder = func(timeout time.Duration) (sshagent.Agent, io.Closer, error) {
return sshagent.NewKeyring(), &testCloser{}, nil
}
routeSSHAgentForwarding = func(client *ssh.Client, keyring sshagent.Agent) error { return nil }
var requestCalls atomic.Int32
requestSSHAgentForwarding = func(session *ssh.Session) error {
@ -241,11 +162,11 @@ func TestNewPTYSessionEnablesAgentForwardingWhenConfigured(t *testing.T) {
func TestNewExecSessionSkipsAgentForwardingWhenDisabled(t *testing.T) {
oldNewSSHSession := newSSHSession
oldProbeSSHAgentForwarding := probeSSHAgentForwarding
oldNewSSHAgentForwarder := newSSHAgentForwarder
oldRequestSSHAgentForwarding := requestSSHAgentForwarding
t.Cleanup(func() {
newSSHSession = oldNewSSHSession
probeSSHAgentForwarding = oldProbeSSHAgentForwarding
newSSHAgentForwarder = oldNewSSHAgentForwarder
requestSSHAgentForwarding = oldRequestSSHAgentForwarding
})
@ -255,9 +176,9 @@ func TestNewExecSessionSkipsAgentForwardingWhenDisabled(t *testing.T) {
newSSHSession = func(client *ssh.Client) (*ssh.Session, error) {
return &ssh.Session{}, nil
}
probeSSHAgentForwarding = func(timeout time.Duration) error {
t.Fatal("agent forwarding probe should not run when disabled")
return nil
newSSHAgentForwarder = func(timeout time.Duration) (sshagent.Agent, io.Closer, error) {
t.Fatal("agent forwarder should not initialize when disabled")
return nil, nil, nil
}
requestSSHAgentForwarding = func(session *ssh.Session) error {
t.Fatal("agent forwarding should not be requested when disabled")
@ -270,18 +191,18 @@ func TestNewExecSessionSkipsAgentForwardingWhenDisabled(t *testing.T) {
}
func TestRequestAgentForwardingReturnsUnavailableError(t *testing.T) {
oldProbeSSHAgentForwarding := probeSSHAgentForwarding
oldNewSSHAgentForwarder := newSSHAgentForwarder
oldRequestSSHAgentForwarding := requestSSHAgentForwarding
t.Cleanup(func() {
probeSSHAgentForwarding = oldProbeSSHAgentForwarding
newSSHAgentForwarder = oldNewSSHAgentForwarder
requestSSHAgentForwarding = oldRequestSSHAgentForwarding
})
star := &StarSSH{}
star.setTransport(&ssh.Client{}, nil)
probeSSHAgentForwarding = func(timeout time.Duration) error {
return errors.New("ssh-agent forwarding unavailable: ssh-agent unavailable")
newSSHAgentForwarder = func(timeout time.Duration) (sshagent.Agent, io.Closer, error) {
return nil, nil, errors.New("ssh-agent forwarding unavailable: ssh-agent unavailable")
}
requestSSHAgentForwarding = func(session *ssh.Session) error {
t.Fatal("session request should not run when agent forwarder init fails")
@ -295,16 +216,16 @@ func TestRequestAgentForwardingReturnsUnavailableError(t *testing.T) {
}
func TestRequestAgentForwardingWrapsSetupErrorAsUnavailable(t *testing.T) {
oldProbeSSHAgentForwarding := probeSSHAgentForwarding
oldNewSSHAgentForwarder := newSSHAgentForwarder
t.Cleanup(func() {
probeSSHAgentForwarding = oldProbeSSHAgentForwarding
newSSHAgentForwarder = oldNewSSHAgentForwarder
})
star := &StarSSH{}
star.setTransport(&ssh.Client{}, nil)
probeSSHAgentForwarding = func(timeout time.Duration) error {
return errors.New("dial unix /tmp/ssh-broken.sock: connect: permission denied")
newSSHAgentForwarder = func(timeout time.Duration) (sshagent.Agent, io.Closer, error) {
return nil, nil, errors.New("dial unix /tmp/ssh-broken.sock: connect: permission denied")
}
err := star.RequestAgentForwarding(&ssh.Session{})
@ -314,11 +235,11 @@ func TestRequestAgentForwardingWrapsSetupErrorAsUnavailable(t *testing.T) {
}
func TestRequestAgentForwardingReturnsDeniedError(t *testing.T) {
oldProbeSSHAgentForwarding := probeSSHAgentForwarding
oldNewSSHAgentForwarder := newSSHAgentForwarder
oldRouteSSHAgentForwarding := routeSSHAgentForwarding
oldRequestSSHAgentForwarding := requestSSHAgentForwarding
t.Cleanup(func() {
probeSSHAgentForwarding = oldProbeSSHAgentForwarding
newSSHAgentForwarder = oldNewSSHAgentForwarder
routeSSHAgentForwarding = oldRouteSSHAgentForwarding
requestSSHAgentForwarding = oldRequestSSHAgentForwarding
})
@ -326,12 +247,10 @@ func TestRequestAgentForwardingReturnsDeniedError(t *testing.T) {
star := &StarSSH{}
star.setTransport(&ssh.Client{}, nil)
probeSSHAgentForwarding = func(timeout time.Duration) error {
return nil
}
routeSSHAgentForwarding = func(client *ssh.Client, timeout time.Duration) (io.Closer, error) {
return &testCloser{}, nil
newSSHAgentForwarder = func(timeout time.Duration) (sshagent.Agent, io.Closer, error) {
return sshagent.NewKeyring(), &testCloser{}, nil
}
routeSSHAgentForwarding = func(client *ssh.Client, keyring sshagent.Agent) error { return nil }
requestSSHAgentForwarding = func(session *ssh.Session) error {
return errors.New("forwarding request denied")
}
@ -344,12 +263,12 @@ func TestRequestAgentForwardingReturnsDeniedError(t *testing.T) {
func TestNewExecSessionIgnoresAgentForwardingDenied(t *testing.T) {
oldNewSSHSession := newSSHSession
oldProbeSSHAgentForwarding := probeSSHAgentForwarding
oldNewSSHAgentForwarder := newSSHAgentForwarder
oldRouteSSHAgentForwarding := routeSSHAgentForwarding
oldRequestSSHAgentForwarding := requestSSHAgentForwarding
t.Cleanup(func() {
newSSHSession = oldNewSSHSession
probeSSHAgentForwarding = oldProbeSSHAgentForwarding
newSSHAgentForwarder = oldNewSSHAgentForwarder
routeSSHAgentForwarding = oldRouteSSHAgentForwarding
requestSSHAgentForwarding = oldRequestSSHAgentForwarding
})
@ -364,12 +283,10 @@ func TestNewExecSessionIgnoresAgentForwardingDenied(t *testing.T) {
newSSHSession = func(client *ssh.Client) (*ssh.Session, error) {
return &ssh.Session{}, nil
}
probeSSHAgentForwarding = func(timeout time.Duration) error {
return nil
}
routeSSHAgentForwarding = func(client *ssh.Client, timeout time.Duration) (io.Closer, error) {
return &testCloser{}, nil
newSSHAgentForwarder = func(timeout time.Duration) (sshagent.Agent, io.Closer, error) {
return sshagent.NewKeyring(), &testCloser{}, nil
}
routeSSHAgentForwarding = func(client *ssh.Client, keyring sshagent.Agent) error { return nil }
requestSSHAgentForwarding = func(session *ssh.Session) error {
return errors.New("forwarding request denied")
}
@ -381,10 +298,10 @@ func TestNewExecSessionIgnoresAgentForwardingDenied(t *testing.T) {
func TestNewExecSessionIgnoresAgentForwardingUnavailable(t *testing.T) {
oldNewSSHSession := newSSHSession
oldProbeSSHAgentForwarding := probeSSHAgentForwarding
oldNewSSHAgentForwarder := newSSHAgentForwarder
t.Cleanup(func() {
newSSHSession = oldNewSSHSession
probeSSHAgentForwarding = oldProbeSSHAgentForwarding
newSSHAgentForwarder = oldNewSSHAgentForwarder
})
star := &StarSSH{
@ -397,8 +314,8 @@ func TestNewExecSessionIgnoresAgentForwardingUnavailable(t *testing.T) {
newSSHSession = func(client *ssh.Client) (*ssh.Session, error) {
return &ssh.Session{}, nil
}
probeSSHAgentForwarding = func(timeout time.Duration) error {
return errors.New("ssh-agent forwarding unavailable: ssh-agent unavailable")
newSSHAgentForwarder = func(timeout time.Duration) (sshagent.Agent, io.Closer, error) {
return nil, nil, errors.New("ssh-agent forwarding unavailable: ssh-agent unavailable")
}
if _, err := star.NewExecSession(); err != nil {
@ -408,10 +325,10 @@ func TestNewExecSessionIgnoresAgentForwardingUnavailable(t *testing.T) {
func TestNewExecSessionIgnoresAgentForwardingSetupError(t *testing.T) {
oldNewSSHSession := newSSHSession
oldProbeSSHAgentForwarding := probeSSHAgentForwarding
oldNewSSHAgentForwarder := newSSHAgentForwarder
t.Cleanup(func() {
newSSHSession = oldNewSSHSession
probeSSHAgentForwarding = oldProbeSSHAgentForwarding
newSSHAgentForwarder = oldNewSSHAgentForwarder
})
star := &StarSSH{
@ -424,8 +341,8 @@ func TestNewExecSessionIgnoresAgentForwardingSetupError(t *testing.T) {
newSSHSession = func(client *ssh.Client) (*ssh.Session, error) {
return &ssh.Session{}, nil
}
probeSSHAgentForwarding = func(timeout time.Duration) error {
return errors.New("dial unix /tmp/ssh-broken.sock: connect: connection refused")
newSSHAgentForwarder = func(timeout time.Duration) (sshagent.Agent, io.Closer, error) {
return nil, nil, errors.New("dial unix /tmp/ssh-broken.sock: connect: connection refused")
}
if _, err := star.NewExecSession(); err != nil {
@ -434,11 +351,11 @@ func TestNewExecSessionIgnoresAgentForwardingSetupError(t *testing.T) {
}
func TestEnsureAgentForwardingClosesNewForwarderWhenCloseStarts(t *testing.T) {
oldProbeSSHAgentForwarding := probeSSHAgentForwarding
oldNewSSHAgentForwarder := newSSHAgentForwarder
oldRouteSSHAgentForwarding := routeSSHAgentForwarding
oldCloseSSHClient := closeSSHClient
t.Cleanup(func() {
probeSSHAgentForwarding = oldProbeSSHAgentForwarding
newSSHAgentForwarder = oldNewSSHAgentForwarder
routeSSHAgentForwarding = oldRouteSSHAgentForwarding
closeSSHClient = oldCloseSSHClient
})
@ -453,13 +370,13 @@ func TestEnsureAgentForwardingClosesNewForwarderWhenCloseStarts(t *testing.T) {
started := make(chan struct{})
release := make(chan struct{})
closer := &testCloser{}
probeSSHAgentForwarding = func(timeout time.Duration) error {
return nil
}
routeSSHAgentForwarding = func(client *ssh.Client, timeout time.Duration) (io.Closer, error) {
newSSHAgentForwarder = func(timeout time.Duration) (sshagent.Agent, io.Closer, error) {
close(started)
<-release
return closer, nil
return sshagent.NewKeyring(), closer, nil
}
routeSSHAgentForwarding = func(client *ssh.Client, keyring sshagent.Agent) error {
return nil
}
closeSSHClient = func(client sshClientRequester) error { return nil }
@ -498,75 +415,3 @@ func TestEnsureAgentForwardingClosesNewForwarderWhenCloseStarts(t *testing.T) {
t.Fatal("expected no leaked agent forwarder after close race")
}
}
func TestProxySSHAgentChannelClosesBlockedAgentConnWhenRemoteChannelEnds(t *testing.T) {
agentConn, peerConn := net.Pipe()
defer peerConn.Close()
tracked := &trackedConn{Conn: agentConn}
channel := newTestSSHChannel(func(p []byte) (int, error) {
return 0, io.EOF
})
done := make(chan struct{})
go func() {
proxySSHAgentChannel(channel, tracked)
close(done)
}()
select {
case <-done:
case <-time.After(time.Second):
t.Fatal("proxySSHAgentChannel did not exit after remote EOF")
}
if tracked.closed.Load() == 0 {
t.Fatal("expected local agent connection to be closed")
}
if channel.closed.Load() == 0 {
t.Fatal("expected ssh channel to be closed")
}
}
func TestSSHAgentForwardProxyCloseClosesActiveBridges(t *testing.T) {
agentConn, peerConn := net.Pipe()
defer peerConn.Close()
tracked := &trackedConn{Conn: agentConn}
channel := newBlockingTestSSHChannel()
proxy := &sshAgentForwardProxy{
stopCh: make(chan struct{}),
active: make(map[*sshAgentForwardBridge]struct{}),
}
bridge := &sshAgentForwardBridge{
proxy: proxy,
channel: channel,
conn: tracked,
}
if !proxy.registerBridge(bridge) {
t.Fatal("expected bridge registration to succeed")
}
done := make(chan struct{})
go func() {
bridge.run()
close(done)
}()
if err := proxy.Close(); err != nil {
t.Fatalf("close proxy: %v", err)
}
select {
case <-done:
case <-time.After(time.Second):
t.Fatal("bridge did not exit after proxy close")
}
if tracked.closed.Load() == 0 {
t.Fatal("expected proxy close to close local agent connection")
}
if channel.closed.Load() == 0 {
t.Fatal("expected proxy close to close ssh channel")
}
}

View File

@ -16,7 +16,6 @@ import (
var ErrHostKeyCallbackRequired = errors.New("host key callback is required; use DefaultAllowHostKeyCallback to explicitly allow any host key")
var errSSHAgentUnavailable = errors.New("ssh-agent unavailable")
var buildSSHAgentAuthMethodFunc = buildSSHAgentAuthMethod
var defaultAuthOrder = []AuthMethodKind{
AuthMethodSSHAgent,
@ -43,8 +42,7 @@ func loginWithContext(ctx context.Context, info LoginInput) (*StarSSH, error) {
return nil, ErrHostKeyCallbackRequired
}
authTimeout := effectiveLoginTimeout(info)
loginCtx, cancel := contextWithLoginTimeout(ctx, authTimeout)
loginCtx, cancel := contextWithLoginTimeout(ctx, info.Timeout)
defer cancel()
sshInfo := &StarSSH{
@ -78,7 +76,7 @@ func loginWithContext(ctx context.Context, info LoginInput) (*StarSSH, error) {
clientConfig := &ssh.ClientConfig{
User: info.User,
Auth: auth,
Timeout: authTimeout,
Timeout: info.Timeout,
HostKeyCallback: hostKeyCallback,
BannerCallback: bannerCallback,
}
@ -95,7 +93,7 @@ func loginWithContext(ctx context.Context, info LoginInput) (*StarSSH, error) {
if err != nil {
return sshInfo, err
}
restoreDeadline := applyConnDeadline(rawConn, loginCtx, authTimeout)
restoreDeadline := applyConnDeadline(rawConn, loginCtx, info.Timeout)
defer restoreDeadline()
clientConn, chans, reqs, err := ssh.NewClientConn(rawConn, targetAddr, clientConfig)
@ -132,7 +130,6 @@ func LoginSimple(host string, user string, passwd string, prikeyPath string, por
Addr: host,
Port: port,
Timeout: timeout,
DialTimeout: timeout,
User: user,
HostKeyCallback: DefaultAllowHostKeyCallback,
}
@ -157,27 +154,10 @@ func normalizeLoginInput(info LoginInput) LoginInput {
if info.Port <= 0 {
info.Port = defaultSSHPort
}
return info
}
func effectiveLoginTimeout(info LoginInput) time.Duration {
if info.Timeout <= 0 {
return 0
}
return info.Timeout
}
func effectiveDialTimeout(info LoginInput) time.Duration {
switch {
case info.DialTimeout < 0:
return 0
case info.DialTimeout > 0:
return info.DialTimeout
case info.Timeout > 0:
return info.Timeout
default:
return defaultLoginTimeout
info.Timeout = defaultLoginTimeout
}
return info
}
func buildAuthMethods(info LoginInput) ([]ssh.AuthMethod, func(), error) {
@ -214,7 +194,7 @@ func buildAuthMethods(info LoginInput) ([]ssh.AuthMethod, func(), error) {
if info.DisableSSHAgent {
continue
}
agentMethod, cleanup, err := buildSSHAgentAuthMethodFunc(effectiveDialTimeout(info))
agentMethod, cleanup, err := buildSSHAgentAuthMethod(info.Timeout)
if err != nil {
agentErr = err
continue

View File

@ -1,99 +0,0 @@
package starssh
import (
"testing"
"time"
"golang.org/x/crypto/ssh"
)
func TestNormalizeLoginInputKeepsZeroAuthTimeout(t *testing.T) {
info := normalizeLoginInput(LoginInput{})
if info.Port != defaultSSHPort {
t.Fatalf("Port=%d want %d", info.Port, defaultSSHPort)
}
if info.Timeout != 0 {
t.Fatalf("Timeout=%v want 0", info.Timeout)
}
if info.DialTimeout != 0 {
t.Fatalf("DialTimeout=%v want 0", info.DialTimeout)
}
}
func TestEffectiveLoginTimeout(t *testing.T) {
if got := effectiveLoginTimeout(LoginInput{}); got != 0 {
t.Fatalf("zero login timeout should stay zero, got %v", got)
}
if got := effectiveLoginTimeout(LoginInput{Timeout: 7 * time.Second}); got != 7*time.Second {
t.Fatalf("expected explicit login timeout, got %v", got)
}
}
func TestEffectiveDialTimeout(t *testing.T) {
tests := []struct {
name string
info LoginInput
want time.Duration
}{
{
name: "default fallback",
info: LoginInput{},
want: defaultLoginTimeout,
},
{
name: "reuse timeout when dial timeout omitted",
info: LoginInput{Timeout: 9 * time.Second},
want: 9 * time.Second,
},
{
name: "explicit dial timeout wins",
info: LoginInput{Timeout: 9 * time.Second, DialTimeout: 3 * time.Second},
want: 3 * time.Second,
},
{
name: "negative dial timeout disables default dial deadline",
info: LoginInput{Timeout: 9 * time.Second, DialTimeout: -1},
want: 0,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
if got := effectiveDialTimeout(tc.info); got != tc.want {
t.Fatalf("effectiveDialTimeout(%+v)=%v want %v", tc.info, got, tc.want)
}
})
}
}
func TestBuildAuthMethodsUsesDialTimeoutInsteadOfAuthTimeout(t *testing.T) {
oldBuilder := buildSSHAgentAuthMethodFunc
t.Cleanup(func() {
buildSSHAgentAuthMethodFunc = oldBuilder
})
captured := time.Duration(-2)
buildSSHAgentAuthMethodFunc = func(timeout time.Duration) (ssh.AuthMethod, func(), error) {
captured = timeout
return ssh.Password("agent"), nil, nil
}
info := LoginInput{
Timeout: 0,
DialTimeout: 11 * time.Second,
AuthOrder: []AuthMethodKind{AuthMethodSSHAgent},
}
auth, cleanup, err := buildAuthMethods(info)
if err != nil {
t.Fatalf("buildAuthMethods: %v", err)
}
if cleanup != nil {
cleanup()
}
if len(auth) != 1 {
t.Fatalf("expected one auth method, got %d", len(auth))
}
if captured != 11*time.Second {
t.Fatalf("agent auth builder timeout=%v want %v", captured, 11*time.Second)
}
}

View File

@ -32,7 +32,7 @@ func resolveDialContext(info LoginInput) DialContextFunc {
}
dialer := &net.Dialer{
Timeout: effectiveDialTimeout(info),
Timeout: info.Timeout,
}
return dialer.DialContext
}
@ -44,7 +44,7 @@ func dialTargetConn(ctx context.Context, info LoginInput) (net.Conn, *StarSSH, e
}
dialContext := resolveDialContext(info)
proxyConfig := normalizeProxyConfig(info.Proxy, effectiveDialTimeout(info))
proxyConfig := normalizeProxyConfig(info.Proxy, info.Timeout)
if proxyConfig != nil {
return dialViaProxy(ctx, dialContext, *proxyConfig, targetAddr)
}

View File

@ -94,20 +94,14 @@ type LoginInput struct {
AuthOrder []AuthMethodKind
Addr string
Port int
// Timeout limits the SSH handshake/authentication phase after a TCP connection has
// already been established. Zero means no authentication timeout.
Timeout time.Duration
// DialTimeout limits outbound dial steps such as TCP connect, proxy connect, and
// local ssh-agent socket connect. Zero falls back to Timeout when set, otherwise
// uses the package default dial timeout. Negative disables the default dial timeout.
DialTimeout time.Duration
DialContext DialContextFunc
Proxy *ProxyConfig
Jump *LoginInput
KeepAliveInterval time.Duration
KeepAliveTimeout time.Duration
HostKeyCallback func(string, net.Addr, ssh.PublicKey) error
BannerCallback func(string) error
Timeout time.Duration
DialContext DialContextFunc
Proxy *ProxyConfig
Jump *LoginInput
KeepAliveInterval time.Duration
KeepAliveTimeout time.Duration
HostKeyCallback func(string, net.Addr, ssh.PublicKey) error
BannerCallback func(string) error
}
// StarShell keeps the legacy prompt-driven helper for POSIX-style scripted shell interactions.