Release 2.6.2: disconnect AWG without freezing the UI.

Co-authored-by: Cursor <cursoragent@cursor.com>

internal/protocols/awg/engine.go

@@ -120,15 +120,16 @@ func (e *Engine) stopLocked() error {
 	if !e.running {
 		return nil
 	}
-	if e.proxies != nil {
-		_ = e.proxies.Close()
-		e.proxies = nil
-	}
+	// Close device first so tun/netstack dials fail and relays exit promptly.
 	if e.dev != nil {
 		e.dev.Close()
 		e.dev = nil
 	}
 	e.tnet = nil
+	if e.proxies != nil {
+		_ = e.proxies.Close()
+		e.proxies = nil
+	}
 	e.running = false
 	return nil
 }

internal/protocols/awg/proxy.go

@@ -17,18 +17,22 @@ import (
 type DialFunc func(ctx context.Context, network, address string) (net.Conn, error)
 
 type proxyServer struct {
-	mu       sync.Mutex
-	socksLn  net.Listener
-	httpLn   net.Listener
-	dial     DialFunc
-	wg       sync.WaitGroup
-	closed   bool
+	mu        sync.Mutex
+	socksLn   net.Listener
+	httpLn    net.Listener
+	dial      DialFunc
+	wg        sync.WaitGroup
+	closed    bool
 	socksAddr string
 	httpAddr  string
+	conns     map[net.Conn]struct{}
 }
 
 func startProxies(socksHostPort, httpHostPort string, dial DialFunc) (*proxyServer, error) {
-	s := &proxyServer{dial: dial}
+	s := &proxyServer{
+		dial:  dial,
+		conns: make(map[net.Conn]struct{}),
+	}
 	if socksHostPort != "" {
 		ln, err := net.Listen("tcp", socksHostPort)
 		if err != nil {
@@ -53,6 +57,22 @@ func startProxies(socksHostPort, httpHostPort string, dial DialFunc) (*proxyServ
 	return s, nil
 }
 
+func (s *proxyServer) track(c net.Conn) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.closed {
+		_ = c.Close()
+		return
+	}
+	s.conns[c] = struct{}{}
+}
+
+func (s *proxyServer) untrack(c net.Conn) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	delete(s.conns, c)
+}
+
 func (s *proxyServer) Close() error {
 	s.mu.Lock()
 	if s.closed {
@@ -66,8 +86,23 @@ func (s *proxyServer) Close() error {
 	if s.httpLn != nil {
 		_ = s.httpLn.Close()
 	}
+	// Force-close active relays so io.Copy / dial wake up immediately.
+	for c := range s.conns {
+		_ = c.Close()
+	}
+	s.conns = nil
 	s.mu.Unlock()
-	s.wg.Wait()
+
+	done := make(chan struct{})
+	go func() {
+		s.wg.Wait()
+		close(done)
+	}()
+	select {
+	case <-done:
+	case <-time.After(1500 * time.Millisecond):
+		// Don't hang UI / Disconnect on stuck netstack copies.
+	}
 	return nil
 }
 
@@ -81,6 +116,8 @@ func (s *proxyServer) serveSOCKS() {
 		s.wg.Add(1)
 		go func(conn net.Conn) {
 			defer s.wg.Done()
+			s.track(conn)
+			defer s.untrack(conn)
 			_ = handleSOCKS(conn, s.dial)
 		}(c)
 	}
@@ -96,6 +133,8 @@ func (s *proxyServer) serveHTTP() {
 		s.wg.Add(1)
 		go func(conn net.Conn) {
 			defer s.wg.Done()
+			s.track(conn)
+			defer s.untrack(conn)
 			_ = handleHTTPProxy(conn, s.dial)
 		}(c)
 	}
@@ -156,7 +195,7 @@ func handleSOCKS(client net.Conn, dial DialFunc) error {
 	port := binary.BigEndian.Uint16(buf[:2])
 	_ = client.SetDeadline(time.Time{})
 
-	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
 	defer cancel()
 	remote, err := dial(ctx, "tcp", net.JoinHostPort(host, fmt.Sprintf("%d", port)))
 	if err != nil {
@@ -185,7 +224,7 @@ func handleHTTPProxy(client net.Conn, dial DialFunc) error {
 		if !strings.Contains(host, ":") {
 			host += ":443"
 		}
-		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+		ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
 		defer cancel()
 		remote, err := dial(ctx, "tcp", host)
 		if err != nil {
@@ -212,7 +251,7 @@ func handleHTTPProxy(client net.Conn, dial DialFunc) error {
 	if !strings.Contains(host, ":") {
 		host += ":80"
 	}
-	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
 	defer cancel()
 	remote, err := dial(ctx, "tcp", host)
 	if err != nil {