diff --git a/tcpproxy.go b/tcpproxy.go index d59c434..4dfd5ab 100644 --- a/tcpproxy.go +++ b/tcpproxy.go @@ -293,6 +293,12 @@ type Target interface { HandleConn(net.Conn) } +type contextKey struct{} + +// SourceAddrContextKey is the context key used by DialProxy.HandleConn to +// pass the incoming connection's remote address (net.Addr) to DialContext. +var SourceAddrContextKey = contextKey{} + // To is shorthand way of writing &tcpproxy.DialProxy{Addr: addr}. func To(addr string) *DialProxy { return &DialProxy{Addr: addr} @@ -374,7 +380,7 @@ func closeWrite(c net.Conn) { // HandleConn implements the Target interface. func (dp *DialProxy) HandleConn(src net.Conn) { - ctx := context.Background() + ctx := context.WithValue(context.Background(), SourceAddrContextKey, src.RemoteAddr()) var cancel context.CancelFunc if dp.DialTimeout >= 0 { ctx, cancel = context.WithTimeout(ctx, dp.dialTimeout())