diff --git a/internal/daemon/worker/proxy/proxydialer.go b/internal/daemon/worker/proxy/proxydialer.go index 8d23a4814..9df3d9895 100644 --- a/internal/daemon/worker/proxy/proxydialer.go +++ b/internal/daemon/worker/proxy/proxydialer.go @@ -2,6 +2,7 @@ package proxy import ( "context" + "fmt" "net" "sync/atomic" @@ -20,11 +21,15 @@ func directDialer(ctx context.Context, endpoint string, _ string, _ proto.Messag if len(endpoint) == 0 { return nil, errors.New(ctx, errors.InvalidParameter, op, "endpoint is empty") } - d, err := NewProxyDialer(ctx, func(...Option) (net.Conn, error) { + d, err := NewProxyDialer(ctx, func(opt ...Option) (net.Conn, error) { remoteConn, err := net.Dial("tcp", endpoint) if err != nil { return nil, errors.Wrap(ctx, err, op) } + opts := GetOpts(opt...) + if opts.WithPostConnectionHook != nil { + opts.WithPostConnectionHook(remoteConn) + } return remoteConn, nil }) if err != nil { @@ -94,16 +99,17 @@ func (d *ProxyDialer) Dial(ctx context.Context, opt ...Option) (net.Conn, error) if err != nil { return nil, err } - switch v := c.RemoteAddr().(type) { - case *net.TCPAddr: - ip := v.IP.String() - port := uint32(v.Port) + switch v := c.(type) { + case *net.TCPConn: + addr := v.RemoteAddr().(*net.TCPAddr) + ip := addr.IP.String() + port := uint32(addr.Port) d.latestAddr.Store(&proxyAddr{ip: ip, port: port}) case portAndIpGetter: d.latestAddr.Store(&proxyAddr{ip: v.GetIp(), port: v.GetPort()}) default: c.Close() - return nil, errors.New(ctx, errors.Internal, op, "connection type unexpected") + return nil, errors.New(ctx, errors.Internal, op, fmt.Sprintf("connection type unexpected %T", v)) } return c, nil } diff --git a/internal/daemon/worker/proxy/proxydialer_test.go b/internal/daemon/worker/proxy/proxydialer_test.go index 9822b6795..d1bb10640 100644 --- a/internal/daemon/worker/proxy/proxydialer_test.go +++ b/internal/daemon/worker/proxy/proxydialer_test.go @@ -60,6 +60,3 @@ func TestProxyDialer(t *testing.T) { assert.EqualValues(t, tcpAddr.Port, d.LastConnectionAddr().Port()) }) } - -func TestProxyDialer_Dial(t *testing.T) { -}