diff --git a/redirector/redirector.go b/redirector/redirector.go index 62aba67..c5910f3 100644 --- a/redirector/redirector.go +++ b/redirector/redirector.go @@ -30,8 +30,12 @@ type Redirector struct { } func (r *Redirector) Redirect(redirection *Redirection) { - r.redirectionChan <- redirection - log.Debug("redirect request") + select { + case r.redirectionChan <- redirection: + log.Debug("redirect request") + case <-r.ctx.Done(): + log.Debug("exiting") + } } func (r *Redirector) worker() { @@ -44,13 +48,13 @@ func (r *Redirector) worker() { return } defer redirection.InboundConn.Close() - if redirection.Dial == nil { - redirection.Dial = defaultDial - } if redirection.RedirectTo == nil || reflect.ValueOf(redirection.RedirectTo).IsNil() { log.Error("nil redirection addr") return } + if redirection.Dial == nil { + redirection.Dial = defaultDial + } log.Warn("redirecting connection from", redirection.InboundConn.RemoteAddr(), "to", redirection.RedirectTo.String()) outboundConn, err := redirection.Dial(redirection.RedirectTo) if err != nil { diff --git a/redirector/redirector_test.go b/redirector/redirector_test.go index 1252e9c..f3151ec 100644 --- a/redirector/redirector_test.go +++ b/redirector/redirector_test.go @@ -35,13 +35,16 @@ func TestRedirector(t *testing.T) { common.Must(err) redir.Redirect(&Redirection{ Dial: nil, - RedirectTo: tunnel.NewAddressFromHostPort("tcp", util.EchoAddr, util.EchoPort), + RedirectTo: tunnel.NewAddressFromHostPort("tcp", "127.0.0.1", util.EchoPort), InboundConn: conn2, }) payload := util.GeneratePayload(1024) common.Must2(conn1.Write(payload)) buf := make([]byte, 1024) - io.ReadFull(conn2, buf) + n, err := io.ReadFull(conn2, buf) + if n != 1024 || err != nil { + t.Fatal(n, err) + } if !bytes.Equal(buf, payload) { t.Fatal("diff: ", payload, "\n", buf) }