From 8a56389396a8eb853777820b35c68dd6137d1128 Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong@windscribe.com>
Date: Mon, 20 Mar 2023 22:53:42 +0700
Subject: [PATCH] cmd/ctrld: ensure both udp/tcp listener aborted

So either one of them return an error, the other will be terminated.
---
 cmd/ctrld/dns_proxy.go | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/cmd/ctrld/dns_proxy.go b/cmd/ctrld/dns_proxy.go
index 733ae63..5143659 100644
--- a/cmd/ctrld/dns_proxy.go
+++ b/cmd/ctrld/dns_proxy.go
@@ -56,7 +56,7 @@ func (p *prog) serveDNS(listenerNum string) error {
 		}
 	})
 
-	g := new(errgroup.Group)
+	g, ctx := errgroup.WithContext(context.Background())
 	for _, proto := range []string{"udp", "tcp"} {
 		proto := proto
 		// On Windows, there's no easy way for disabling/removing IPv6 DNS resolver, so we check whether we can
@@ -68,6 +68,10 @@ func (p *prog) serveDNS(listenerNum string) error {
 					Net:     proto,
 					Handler: handler,
 				}
+				go func() {
+					<-ctx.Done()
+					_ = s.Shutdown()
+				}()
 				if err := s.ListenAndServe(); err != nil {
 					mainLog.Error().Err(err).Msg("could not serving on ::1")
 				}
@@ -80,6 +84,10 @@ func (p *prog) serveDNS(listenerNum string) error {
 				Net:     proto,
 				Handler: handler,
 			}
+			go func() {
+				<-ctx.Done()
+				_ = s.Shutdown()
+			}()
 			if err := s.ListenAndServe(); err != nil {
 				mainLog.Error().Err(err).Msgf("could not listen and serve on: %s", s.Addr)
 				return err