fix tuna exit not started in udp only mode

nknorg · Mar 6, 2023 · 7bd69ed · 7bd69ed
1 parent d5639a4
commit 7bd69ed
Showing 1 changed file with 73 additions and 52 deletions.
diff --git a/client.go b/client.go
@@ -48,6 +48,7 @@ type TunaSessionClient struct {
 	onConnect     chan struct{}
 	onClose       chan struct{}
 	connectedOnce sync.Once
+	udpConn       *udp.EncryptUDPConn
 
 	sync.RWMutex
 	listeners        []net.Listener
@@ -157,69 +158,25 @@ func (c *TunaSessionClient) newTunaExit(i int) (*tuna.TunaExit, error) {
 }
 
 func (c *TunaSessionClient) Listen(addrsRe *nkngomobile.StringArray) error {
-	acceptAddrs, err := getAcceptAddrs(addrsRe)
-	if err != nil {
-		return err
-	}
-
 	c.Lock()
 	defer c.Unlock()
 
-	c.acceptAddrs = acceptAddrs
-
 	if len(c.listeners) > 0 {
 		return nil
 	}
 
-	if c.wallet == nil {
-		return errors.New("wallet is empty")
-	}
-
-	listeners := make([]net.Listener, c.config.NumTunaListeners)
-	for i := 0; i < len(listeners); i++ {
-		listeners[i], err = net.Listen("tcp", "127.0.0.1:")
-		if err != nil {
-			return err
-		}
+	acceptAddrs, err := getAcceptAddrs(addrsRe)
+	if err != nil {
+		return err
 	}
-	c.listeners = listeners
-
-	exits := make([]*tuna.TunaExit, c.config.NumTunaListeners)
-	for i := 0; i < len(listeners); i++ {
-		exits[i], err = c.newTunaExit(i)
-		if err != nil {
-			return err
-		}
 
-		go func(te *tuna.TunaExit) {
-			select {
-			case <-te.OnConnect.C:
-				c.connectedOnce.Do(func() {
-					close(c.onConnect)
-				})
-			case <-c.onClose:
-				return
-			}
-		}(exits[i])
+	c.acceptAddrs = acceptAddrs
 
-		go exits[i].StartReverse(true)
+	if c.wallet == nil {
+		return errors.New("wallet is empty")
 	}
 
-	c.tunaExits = exits
-
-	go func() {
-		select {
-		case <-c.onConnect:
-			go c.listenNKN()
-			for i := 0; i < len(listeners); i++ {
-				go c.listenNet(i)
-			}
-		case <-c.onClose:
-			return
-		}
-	}()
-
-	return nil
+	return c.startExits()
 }
 
 // OnConnect returns a channel that will be closed when at least one tuna exit
@@ -842,17 +799,79 @@ func (c *TunaSessionClient) removeClosedSessions() {
 	}
 }
 
+func (c *TunaSessionClient) startExits() error {
+	if len(c.tunaExits) > 0 {
+		return nil
+	}
+	listeners := make([]net.Listener, c.config.NumTunaListeners)
+	var err error
+	for i := 0; i < len(listeners); i++ {
+		listeners[i], err = net.Listen("tcp", "127.0.0.1:")
+		if err != nil {
+			return err
+		}
+	}
+	c.listeners = listeners
+
+	exits := make([]*tuna.TunaExit, c.config.NumTunaListeners)
+	for i := 0; i < len(listeners); i++ {
+		exits[i], err = c.newTunaExit(i)
+		if err != nil {
+			return err
+		}
+
+		go func(te *tuna.TunaExit) {
+			select {
+			case <-te.OnConnect.C:
+				c.connectedOnce.Do(func() {
+					close(c.onConnect)
+				})
+			case <-c.onClose:
+				return
+			}
+		}(exits[i])
+
+		go exits[i].StartReverse(true)
+	}
+	c.tunaExits = exits
+	go func() {
+		select {
+		case <-c.onConnect:
+			go c.listenNKN()
+			for i := 0; i < len(listeners); i++ {
+				go c.listenNet(i)
+			}
+		case <-c.onClose:
+			return
+		}
+	}()
+
+	return nil
+}
+
 func (c *TunaSessionClient) ListenUDP(addrsRe *nkngomobile.StringArray) (*udp.EncryptUDPConn, error) {
+	c.Lock()
+	defer c.Unlock()
+
 	acceptAddrs, err := getAcceptAddrs(addrsRe)
 	if err != nil {
 		return nil, err
 	}
 	c.acceptAddrs = acceptAddrs
 
+	if c.udpConn != nil {
+		return c.udpConn, nil
+	}
+
 	if c.wallet == nil {
 		return nil, errors.New("wallet is empty")
 	}
 
+	err = c.startExits()
+	if err != nil {
+		return nil, err
+	}
+
 	host, portStr, err := net.SplitHostPort(c.listeners[0].Addr().String())
 	if err != nil {
 		return nil, err
@@ -866,7 +885,9 @@ func (c *TunaSessionClient) ListenUDP(addrsRe *nkngomobile.StringArray) (*udp.En
 	if err != nil {
 		return nil, err
 	}
-	return udp.NewEncryptUDPConn(conn), nil
+
+	c.udpConn = udp.NewEncryptUDPConn(conn)
+	return c.udpConn, nil
 }
 
 func (c *TunaSessionClient) DialUDP(remoteAddr string) (*udp.EncryptUDPConn, error) {