diff --git a/conn/bind_stream.go b/conn/bind_stream.go index 62b11ed..99c1297 100644 --- a/conn/bind_stream.go +++ b/conn/bind_stream.go @@ -92,39 +92,9 @@ func (b *BindStream) readStream(ep *streamEndpoint) { } } -func (b *BindStream) listen(port uint16) { +func (b *BindStream) accept(listener net.Listener) { defer b.wg.Done() - for { - select { - case <-b.ctx.Done(): - return - default: - } - - listenDone := make(chan struct{}) - - listener, err := b.listenConfig.Listen(b.ctx, "tcp", ":"+strconv.Itoa(int(port))) - if err != nil { - continue - } - - b.wg.Add(1) - go b.accept(listener, listenDone) - - select { - case <-b.ctx.Done(): - case <-listenDone: - } - - listener.Close() - } -} - -func (b *BindStream) accept(listener net.Listener, listenDone chan struct{}) { - defer b.wg.Done() - defer close(listenDone) - for { select { case <-b.ctx.Done(): @@ -140,11 +110,11 @@ func (b *BindStream) accept(listener net.Listener, listenDone chan struct{}) { conn = b.upgradeConn(conn) b.wg.Add(1) - go b.handleAccepted(conn, listenDone) + go b.handleAccepted(conn) } } -func (b *BindStream) handleAccepted(conn net.Conn, listenDone chan struct{}) { +func (b *BindStream) handleAccepted(conn net.Conn) { defer b.wg.Done() ep, err := streamEndpointFromConn(conn) @@ -155,11 +125,7 @@ func (b *BindStream) handleAccepted(conn net.Conn, listenDone chan struct{}) { b.wg.Add(1) go b.readStream(ep) - select { - case <-b.ctx.Done(): - case <-listenDone: - } - + <-b.ctx.Done() ep.Close() } @@ -205,18 +171,34 @@ func (b *BindStream) upgradeConn(conn net.Conn) net.Conn { return conn } -func (b *BindStream) Open(port uint16) (fns []ReceiveFunc, actualPort uint16, err error) { - b.port = port - +func (b *BindStream) openConnections() error { b.ctx, b.cancel = context.WithCancel(context.Background()) - b.queue = make(chan *streamPacketQueue, 1024) - if port != 0 { + if b.port != 0 { + listener, err := b.listenConfig.Listen(b.ctx, "tcp", ":"+strconv.Itoa(int(b.port))) + if err != nil { + return err + } + b.wg.Add(1) - go b.listen(port) + go func() { + defer b.wg.Done() + <-b.ctx.Done() + listener.Close() + }() + + b.wg.Add(1) + go b.accept(listener) } - return []ReceiveFunc{b.readFaucet()}, port, nil + return nil +} + +func (b *BindStream) Open(port uint16) (fns []ReceiveFunc, actualPort uint16, err error) { + b.queue = make(chan *streamPacketQueue, 1024) + b.port = port + + return []ReceiveFunc{b.readFaucet()}, b.port, b.openConnections() } func (b *BindStream) Send(bufs [][]byte, ep Endpoint) error { @@ -249,12 +231,15 @@ func (b *BindStream) ParseEndpoint(s string) (Endpoint, error) { return streamEndpointFromAddr(s) } -func (b *BindStream) Close() error { +func (b *BindStream) closeConnections() { if b.cancel != nil { b.cancel() } - b.wg.Wait() +} + +func (b *BindStream) Close() error { + b.closeConnections() if b.queue != nil { close(b.queue) diff --git a/conn/mark_unix.go b/conn/mark_unix.go index 06597a3..a2e8d7c 100644 --- a/conn/mark_unix.go +++ b/conn/mark_unix.go @@ -8,7 +8,6 @@ package conn import ( - "context" "runtime" "syscall" @@ -67,10 +66,7 @@ func (s *StdNetBind) SetMark(mark uint32) error { } func (b *BindStream) SetMark(mark uint32) error { - if b.cancel != nil { - b.cancel() - } - b.wg.Wait() + b.closeConnections() ctrl := b.dialer.Control b.dialer.Control = func(network, address string, c syscall.RawConn) error { @@ -100,11 +96,5 @@ func (b *BindStream) SetMark(mark uint32) error { return err } - b.ctx, b.cancel = context.WithCancel(context.Background()) - if b.port != 0 { - b.wg.Add(1) - go b.listen(b.port) - } - - return nil + return b.openConnections() }