inp2p/cmd/inp2ps/main.go

// inp2ps — INP2P Signaling Server
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"syscall"

	"github.com/openp2p-cn/inp2p/internal/server"
	"github.com/openp2p-cn/inp2p/pkg/auth"
	"github.com/openp2p-cn/inp2p/pkg/config"
	"github.com/openp2p-cn/inp2p/pkg/nat"
	"github.com/openp2p-cn/inp2p/pkg/protocol"
)

func main() {
	cfg := config.DefaultServerConfig()

	flag.IntVar(&cfg.WSPort, "ws-port", cfg.WSPort, "WebSocket signaling port")
	flag.IntVar(&cfg.WebPort, "web-port", cfg.WebPort, "Web console port")
	flag.IntVar(&cfg.STUNUDP1, "stun-udp1", cfg.STUNUDP1, "UDP STUN port 1")
	flag.IntVar(&cfg.STUNUDP2, "stun-udp2", cfg.STUNUDP2, "UDP STUN port 2")
	flag.IntVar(&cfg.STUNTCP1, "stun-tcp1", cfg.STUNTCP1, "TCP STUN port 1")
	flag.IntVar(&cfg.STUNTCP2, "stun-tcp2", cfg.STUNTCP2, "TCP STUN port 2")
	flag.StringVar(&cfg.DBPath, "db", cfg.DBPath, "SQLite database path")
	flag.StringVar(&cfg.CertFile, "cert", "", "TLS certificate file")
	flag.StringVar(&cfg.KeyFile, "key", "", "TLS key file")
	flag.IntVar(&cfg.LogLevel, "log-level", cfg.LogLevel, "Log level (0=debug 1=info 2=warn 3=error)")
	token := flag.Uint64("token", 0, "Master authentication token (uint64)")
	user := flag.String("user", "", "Username for token generation (requires -password)")
	pass := flag.String("password", "", "Password for token generation")
	version := flag.Bool("version", false, "Print version and exit")
	flag.Parse()

	if *version {
		fmt.Printf("inp2ps version %s\ncommit: %s\nbuild: %s\ngo: %s\n",
			config.Version, config.GitCommit, config.BuildTime, config.GoVersion)
		os.Exit(0)
	}

	// Token: either direct value or generated from user+password
	if *token > 0 {
		cfg.Token = *token
	} else if *user != "" && *pass != "" {
		cfg.Token = auth.MakeToken(*user, *pass)
		log.Printf("[main] token generated from credentials: %d", cfg.Token)
	}

	cfg.FillFromEnv()

	if err := cfg.Validate(); err != nil {
		log.Fatalf("[main] config error: %v", err)
	}

	log.Printf("[main] inp2ps v%s starting", config.Version)
	log.Printf("[main] WSS :%d | STUN UDP :%d,%d | STUN TCP :%d,%d",
		cfg.WSPort, cfg.STUNUDP1, cfg.STUNUDP2, cfg.STUNTCP1, cfg.STUNTCP2)

	// ─── STUN Servers ───
	stunQuit := make(chan struct{})

	startSTUN := func(proto string, port int, fn func(int, <-chan struct{}) error) {
		go func() {
			log.Printf("[main] %s STUN listening on :%d", proto, port)
			if err := fn(port, stunQuit); err != nil {
				log.Printf("[main] %s STUN :%d error: %v", proto, port, err)
			}
		}()
	}

	startSTUN("UDP", cfg.STUNUDP1, nat.ServeUDPSTUN)
	if cfg.STUNUDP2 != cfg.STUNUDP1 {
		startSTUN("UDP", cfg.STUNUDP2, nat.ServeUDPSTUN)
	}
	startSTUN("TCP", cfg.STUNTCP1, nat.ServeTCPSTUN)
	if cfg.STUNTCP2 != cfg.STUNTCP1 {
		startSTUN("TCP", cfg.STUNTCP2, nat.ServeTCPSTUN)
	}

	// ─── Signaling Server ───
	srv := server.New(cfg)
	srv.StartCleanup()

	mux := http.NewServeMux()
	mux.HandleFunc("/ws", srv.HandleWS)
	mux.HandleFunc("/api/v1/health", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintf(w, `{"status":"ok","version":"%s","nodes":%d}`, config.Version, len(srv.GetOnlineNodes()))
	})
	mux.HandleFunc("/api/v1/sdwans", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(srv.GetSDWAN())
	})
	mux.HandleFunc("/api/v1/sdwan/edit", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		var req protocol.SDWANConfig
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
		if err := srv.SetSDWAN(req); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]any{"error": 0, "message": "ok"})
	})

	// ─── HTTP Listener ───
	ln, err := net.Listen("tcp", fmt.Sprintf(":%d", cfg.WSPort))
	if err != nil {
		log.Fatalf("[main] listen :%d: %v", cfg.WSPort, err)
	}
	log.Printf("[main] signaling server on :%d (no TLS — use reverse proxy for production)", cfg.WSPort)

	httpSrv := &http.Server{Handler: mux}
	go func() {
		if err := httpSrv.Serve(ln); err != http.ErrServerClosed {
			log.Fatalf("[main] serve: %v", err)
		}
	}()

	// ─── Graceful Shutdown ───
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	<-sigCh

	log.Println("[main] shutting down...")
	close(stunQuit)
	srv.Stop()
	httpSrv.Shutdown(context.Background())
	log.Println("[main] goodbye")
}