dns: merge Start and Wait into ListenAndServe

3 files changed, 14 insertions, 22 deletions

diff --git a/lib/dns/dns_test.go b/lib/dns/dns_test.go
index 20f697a7..c88bf491 100644
--- a/lib/dns/dns_test.go
+++ b/lib/dns/dns_test.go
@@ -45,14 +45,12 @@ func TestMain(m *testing.M) {
 
 	_testServer.LoadMasterFile("testdata/kilabit.info")
 
-	_testServer.Start()
+	go _testServer.ListenAndServe()
 
 	// Wait for all listeners running.
 	time.Sleep(500 * time.Millisecond)
 
-	s := m.Run()
-
-	os.Exit(s)
+	os.Exit(m.Run())
 }
 
 func TestQueryType(t *testing.T) {
diff --git a/lib/dns/example_server_test.go b/lib/dns/example_server_test.go
index 9daa0749..96aac8bb 100644
--- a/lib/dns/example_server_test.go
+++ b/lib/dns/example_server_test.go
@@ -52,7 +52,7 @@ func ExampleServer() {
 	// Load records to be served from master file.
 	server.LoadMasterFile("testdata/kilabit.info")
 
-	server.Start()
+	go server.ListenAndServe()
 
 	// Wait for all listeners running.
 	time.Sleep(500 * time.Millisecond)
diff --git a/lib/dns/server.go b/lib/dns/server.go
index 60a4a3cc..0623ee56 100644
--- a/lib/dns/server.go
+++ b/lib/dns/server.go
@@ -323,17 +323,21 @@ func (srv *Server) RestartForwarders(nameServers, fallbackNS []string) {
 }
 
 //
-// Start the server, listening and serve query from clients.
+// ListenAndServe start listening and serve queries from clients.
 //
-func (srv *Server) Start() {
+func (srv *Server) ListenAndServe() (err error) {
 	srv.startAllForwarders()
 
 	go srv.processRequest()
-
 	go srv.serveDoT()
 	go srv.serveDoH()
 	go srv.serveTCP()
 	go srv.serveUDP()
+
+	err = <-srv.errListener
+	srv.Stop()
+
+	return err
 }
 
 //
@@ -348,6 +352,10 @@ func (srv *Server) Stop() {
 	if err != nil {
 		log.Println("dns: error when closing TCP: " + err.Error())
 	}
+	err = srv.dot.Close()
+	if err != nil {
+		log.Println("dns: error when closing DoT: " + err.Error())
+	}
 	err = srv.doh.Close()
 	if err != nil {
 		log.Println("dns: error when closing DoH: " + err.Error())
@@ -357,20 +365,6 @@ func (srv *Server) Stop() {
 }
 
 //
-// Wait for server to be Stop()-ed or when one of listener throw an error.
-//
-func (srv *Server) Wait() (err error) {
-	err = <-srv.errListener
-	if err != nil && err != io.EOF {
-		log.Println(err)
-	}
-
-	srv.Stop()
-
-	return err
-}
-
-//
 // serveDoH listen for request over HTTPS using certificate and key
 // file in parameter.  The path to request is static "/dns-query".
 //