fix: notui

fumiama · Apr 20, 2024 · 77dc207 · 77dc207
1 parent f79b3a3
commit 77dc207
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/main.go b/main.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"os"
 	"runtime"
+	"sync"
 	"time"
 
 	"github.com/fumiama/terasu/dns"
@@ -57,7 +58,10 @@ func main() {
 	}
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
+	wg := sync.WaitGroup{}
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		if *dnsf != "" {
 			f, err := os.Open(*dnsf)
 			if err != nil {
@@ -99,5 +103,9 @@ func main() {
 			logrus.Warnln("download canceled")
 		}
 	}()
-	sc.show(time.Second)
+	if notui {
+		wg.Wait()
+	} else {
+		sc.show(time.Second)
+	}
 }