refactor nat

cmd/dashboard/rpc/rpc.go

@@ -1,10 +1,16 @@
 package rpc
 
 import (
+	"fmt"
+	"net/http"
+	"time"
+
 	"google.golang.org/grpc"
 
+	"github.com/hashicorp/go-uuid"
 	"github.com/naiba/nezha/model"
-	pb "github.com/naiba/nezha/proto"
+	"github.com/naiba/nezha/pkg/utils"
+	"github.com/naiba/nezha/proto"
 	rpcService "github.com/naiba/nezha/service/rpc"
 	"github.com/naiba/nezha/service/singleton"
 )
@@ -12,7 +18,7 @@ import (
 func ServeRPC() *grpc.Server {
 	server := grpc.NewServer()
 	rpcService.NezhaHandlerSingleton = rpcService.NewNezhaHandler()
-	pb.RegisterNezhaServiceServer(server, rpcService.NezhaHandlerSingleton)
+	proto.RegisterNezhaServiceServer(server, rpcService.NezhaHandlerSingleton)
 	return server
 }
 
@@ -69,7 +75,62 @@ func DispatchKeepalive() {
 				continue
 			}
 
-			singleton.SortedServerList[i].TaskStream.Send(&pb.Task{Type: model.TaskTypeKeepalive})
+			singleton.SortedServerList[i].TaskStream.Send(&proto.Task{Type: model.TaskTypeKeepalive})
 		}
 	})
 }
+
+func ServeNAT(w http.ResponseWriter, r *http.Request, natConfig *model.NAT) {
+	singleton.ServerLock.RLock()
+	server := singleton.ServerList[natConfig.ServerID]
+	singleton.ServerLock.RUnlock()
+	if server == nil || server.TaskStream == nil {
+		w.WriteHeader(http.StatusServiceUnavailable)
+		w.Write([]byte("server not found or not connected"))
+		return
+	}
+
+	streamId, err := uuid.GenerateUUID()
+	if err != nil {
+		w.WriteHeader(http.StatusServiceUnavailable)
+		w.Write([]byte(fmt.Sprintf("stream id error: %v", err)))
+		return
+	}
+
+	rpcService.NezhaHandlerSingleton.CreateStream(streamId)
+	defer rpcService.NezhaHandlerSingleton.CloseStream(streamId)
+
+	taskData, err := utils.Json.Marshal(model.TaskNAT{
+		StreamID: streamId,
+		Host:     natConfig.Host,
+	})
+	if err != nil {
+		w.WriteHeader(http.StatusServiceUnavailable)
+		w.Write([]byte(fmt.Sprintf("task data error: %v", err)))
+		return
+	}
+
+	if err := server.TaskStream.Send(&proto.Task{
+		Type: model.TaskTypeNAT,
+		Data: string(taskData),
+	}); err != nil {
+		w.WriteHeader(http.StatusServiceUnavailable)
+		w.Write([]byte(fmt.Sprintf("send task error: %v", err)))
+		return
+	}
+
+	wWrapped, err := utils.NewRequestWrapper(r, w)
+	if err != nil {
+		w.WriteHeader(http.StatusServiceUnavailable)
+		w.Write([]byte(fmt.Sprintf("request wrapper error: %v", err)))
+		return
+	}
+
+	if err := rpcService.NezhaHandlerSingleton.UserConnected(streamId, wWrapped); err != nil {
+		w.WriteHeader(http.StatusServiceUnavailable)
+		w.Write([]byte(fmt.Sprintf("user connected error: %v", err)))
+		return
+	}
+
+	rpcService.NezhaHandlerSingleton.StartStream(streamId, time.Second*10)
+}