From c533dd34abffcbb92dcb59b6dc0376d9b4860d1b Mon Sep 17 00:00:00 2001
From: sjlleo <sjlleo001@gmail.com>
Date: Mon, 20 Jun 2022 22:17:38 +0800
Subject: [PATCH] fix: fast trace crash

---
 fast_trace/fast_trace.go | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/fast_trace/fast_trace.go b/fast_trace/fast_trace.go
index e086e0b..91f2ea2 100644
--- a/fast_trace/fast_trace.go
+++ b/fast_trace/fast_trace.go
@@ -4,11 +4,14 @@ import (
 	"fmt"
 	"log"
 	"net"
+	"os"
+	"os/signal"
 	"time"
 
 	"github.com/xgadget-lab/nexttrace/ipgeo"
 	"github.com/xgadget-lab/nexttrace/printer"
 	"github.com/xgadget-lab/nexttrace/trace"
+	"github.com/xgadget-lab/nexttrace/wshandle"
 )
 
 type FastTracer struct {
@@ -99,6 +102,14 @@ func FastTest(tm bool) {
 
 	ft := FastTracer{}
 
+	// 建立 WebSocket 连接
+	w := wshandle.New()
+	w.Interrupt = make(chan os.Signal, 1)
+	signal.Notify(w.Interrupt, os.Interrupt)
+	defer func() {
+		w.Conn.Close()
+	}()
+
 	if !tm {
 		ft.TracerouteMethod = trace.ICMPTrace
 		fmt.Println("您将默认使用ICMP协议进行路由跟踪，如果您想使用TCP SYN进行路由跟踪，可以加入 -T 参数")