Parcourir la source

fix command listener

Darien Raymond il y a 7 ans
Parent
commit
c1fc7c738a
2 fichiers modifiés avec 20 ajouts et 5 suppressions
  1. 2 0
      app/commander/commander.go
  2. 18 5
      app/commander/outbound.go

+ 2 - 0
app/commander/commander.go

@@ -10,6 +10,7 @@ import (
 	"google.golang.org/grpc"
 	"v2ray.com/core"
 	"v2ray.com/core/common"
+	"v2ray.com/core/common/signal"
 )
 
 type Commander struct {
@@ -58,6 +59,7 @@ func (c *Commander) Start() error {
 
 	listener := &OutboundListener{
 		buffer: make(chan net.Conn, 4),
+		done:   signal.NewDone(),
 	}
 
 	go func() {

+ 18 - 5
app/commander/outbound.go

@@ -11,26 +11,39 @@ import (
 
 type OutboundListener struct {
 	buffer chan net.Conn
+	done   *signal.Done
 }
 
 func (l *OutboundListener) add(conn net.Conn) {
 	select {
 	case l.buffer <- conn:
+	case <-l.done.C():
+		conn.Close()
 	default:
 		conn.Close()
 	}
 }
 
 func (l *OutboundListener) Accept() (net.Conn, error) {
-	c, open := <-l.buffer
-	if !open {
-		return nil, newError("listener closed")
+	select {
+	case <-l.done.C():
+		return nil, newError("listern closed")
+	case c := <-l.buffer:
+		return c, nil
 	}
-	return c, nil
 }
 
 func (l *OutboundListener) Close() error {
-	close(l.buffer)
+	l.done.Close()
+L:
+	for {
+		select {
+		case c := <-l.buffer:
+			c.Close()
+		default:
+			break L
+		}
+	}
 	return nil
 }