133 lines
3.6 KiB
Go
133 lines
3.6 KiB
Go
package cmd
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log"
|
|
"net/http"
|
|
"os"
|
|
"os/signal"
|
|
"sync"
|
|
"syscall"
|
|
"time"
|
|
|
|
_ "github.com/hiddify/hiddify-core/extension/repository"
|
|
"github.com/hiddify/hiddify-core/utils"
|
|
v2 "github.com/hiddify/hiddify-core/v2"
|
|
"github.com/improbable-eng/grpc-web/go/grpcweb"
|
|
"github.com/spf13/cobra"
|
|
"google.golang.org/grpc"
|
|
)
|
|
|
|
var extension_id string
|
|
|
|
var commandExtension = &cobra.Command{
|
|
Use: "extension",
|
|
Short: "extension configuration",
|
|
Args: cobra.MaximumNArgs(2),
|
|
Run: func(cmd *cobra.Command, args []string) {
|
|
StartExtension()
|
|
},
|
|
}
|
|
|
|
func StartExtension() {
|
|
v2.Setup("./tmp", "./", "./tmp", 0, false)
|
|
grpc_server, _ := v2.StartCoreGrpcServer("127.0.0.1:12345")
|
|
fmt.Printf("Waiting for CTRL+C to stop\n")
|
|
runWebserver(grpc_server)
|
|
<-time.After(1 * time.Second)
|
|
}
|
|
|
|
func init() {
|
|
// commandWarp.Flags().StringVarP(&warpKey, "key", "k", "", "warp key")
|
|
mainCommand.AddCommand(commandExtension)
|
|
}
|
|
|
|
func allowCors(resp http.ResponseWriter, req *http.Request) {
|
|
resp.Header().Set("Access-Control-Allow-Origin", "*")
|
|
resp.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
|
|
resp.Header().Set("Access-Control-Allow-Headers", "Content-Type")
|
|
if req.Method == "OPTIONS" {
|
|
resp.WriteHeader(http.StatusOK)
|
|
return
|
|
}
|
|
}
|
|
|
|
func runWebserver(grpcServer *grpc.Server) {
|
|
// Context for cancellation
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
// Channels to signal termination
|
|
grpcTerminated := make(chan struct{})
|
|
grpcWebTerminated := make(chan struct{})
|
|
|
|
// Specify the directory to serve static files
|
|
dir := "./extension/html/"
|
|
|
|
// Wrapping gRPC server with grpc-web
|
|
grpcWeb := grpcweb.WrapServer(grpcServer)
|
|
|
|
// HTTP multiplexer
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("/", func(resp http.ResponseWriter, req *http.Request) {
|
|
allowCors(resp, req)
|
|
if grpcWeb.IsGrpcWebRequest(req) || grpcWeb.IsAcceptableGrpcCorsRequest(req) {
|
|
grpcWeb.ServeHTTP(resp, req)
|
|
} else {
|
|
http.DefaultServeMux.ServeHTTP(resp, req)
|
|
}
|
|
})
|
|
|
|
// File server for static files
|
|
fs := http.FileServer(http.Dir(dir))
|
|
http.Handle("/", http.StripPrefix("/", fs))
|
|
|
|
// HTTP server for grpc-web
|
|
rpcWebServer := &http.Server{
|
|
Handler: mux,
|
|
Addr: ":12346",
|
|
}
|
|
log.Println("Serving grpc-web from https://localhost:12346/")
|
|
|
|
// Add a goroutine for the grpc-web server
|
|
wg := sync.WaitGroup{}
|
|
wg.Add(1)
|
|
|
|
go func() {
|
|
defer wg.Done()
|
|
utils.GenerateCertificate("cert/server-cert.pem", "cert/server-key.pem", true, true)
|
|
if err := rpcWebServer.ListenAndServeTLS("cert/server-cert.pem", "cert/server-key.pem"); err != nil && err != http.ErrServerClosed {
|
|
// if err := rpcWebServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
|
fmt.Printf("Web server (gRPC-web) shutdown with error: %s", err)
|
|
}
|
|
grpcServer.Stop()
|
|
close(grpcWebTerminated) // Server terminated
|
|
}()
|
|
|
|
// Signal handling to gracefully shutdown
|
|
sigChan := make(chan os.Signal, 1)
|
|
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
|
|
|
|
select {
|
|
case <-ctx.Done(): // Context canceled
|
|
log.Println("Context canceled, shutting down servers...")
|
|
case sig := <-sigChan: // OS signal received
|
|
log.Printf("Received signal: %s, shutting down servers...", sig)
|
|
case <-grpcTerminated: // Unexpected gRPC termination
|
|
log.Println("gRPC server terminated unexpectedly")
|
|
case <-grpcWebTerminated: // Unexpected gRPC-web termination
|
|
log.Println("gRPC-web server terminated unexpectedly")
|
|
}
|
|
|
|
// Graceful shutdown of the servers
|
|
if err := rpcWebServer.Shutdown(ctx); err != nil {
|
|
log.Printf("gRPC-web server shutdown with error: %s", err)
|
|
}
|
|
<-grpcWebTerminated
|
|
|
|
// Ensure all routines finish
|
|
wg.Wait()
|
|
log.Println("Server shutdown complete")
|
|
}
|