speedtest/cmd/netstable/main.go

package main

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

	ipdata "speedtest/data"
	appclient "speedtest/internal/client"
	"speedtest/internal/geo"
	"speedtest/internal/probe"
	recordstore "speedtest/internal/store"
	appweb "speedtest/internal/web"
	webassets "speedtest/web"
)

func main() {
	args := os.Args[1:]
	if len(args) > 0 {
		switch args[0] {
		case "client":
			runClient(args[1:])
			return
		case "serve":
			runServer(args[1:])
			return
		}
	}
	runServer(args)
}

func runServer(args []string) {
	flags := flag.NewFlagSet("serve", flag.ExitOnError)
	listen := flags.String("listen", ":8080", "HTTP listen address")
	dataPath := flags.String("data", "data/records.jsonl", "JSONL record file")
	staticDir := flags.String("static", "", "optional static asset directory; embedded assets are used when empty")
	downloadsDir := flags.String("downloads-dir", "", "optional directory for public client binaries, served from /downloads/")
	geoBase := flags.String("geo-base", "https://ipapi.co", "online IP geolocation fallback API base URL; set empty to disable online fallback")
	geoTimeout := flags.Duration("geo-timeout", 2*time.Second, "IP geolocation request timeout")
	trustProxyHeaders := flags.Bool("trust-proxy-headers", false, "trust CF-Connecting-IP, X-Real-IP, and X-Forwarded-For headers")
	bandwidthLimitMbps := flags.Float64("bandwidth-limit-mbps", 0, "theoretical program bandwidth limit in Mbps; 0 means unlimited")
	iperf3Path := flags.String("iperf3-path", "iperf3", "iperf3 executable path for one-off CLI iperf3 sessions")
	_ = flags.Parse(args)

	store := recordstore.New(*dataPath)
	var resolvers []appweb.ClientResolver
	if db, err := ipdata.FileSystem().ReadFile("ip2region_v4.xdb"); err != nil {
		log.Printf("embedded ip2region database unavailable: %v", err)
	} else if localResolver, err := geo.NewXDBResolver(db); err != nil {
		log.Printf("embedded ip2region database failed to load: %v", err)
	} else {
		resolvers = append(resolvers, localResolver.Resolve)
	}
	if *geoBase != "" {
		if *geoBase == "https://ipapi.co" {
			resolvers = append(resolvers, geo.Chain(
				geo.New(*geoBase, *geoTimeout),
				geo.New("https://ipinfo.io", *geoTimeout),
			))
		} else {
			resolvers = append(resolvers, geo.New(*geoBase, *geoTimeout).Resolve)
		}
	}
	resolver := chainResolvers(resolvers...)

	options := appweb.Options{
		Store:              store,
		Runner:             appweb.DefaultRunner,
		ClientResolver:     resolver,
		TrustProxyHeaders:  *trustProxyHeaders,
		BandwidthLimitMbps: *bandwidthLimitMbps,
		Iperf3Path:         *iperf3Path,
		DownloadsDir:       *downloadsDir,
		StaticDir:          *staticDir,
	}
	if *staticDir == "" {
		options.StaticFS = webassets.FileSystem()
	}

	server := appweb.New(options)
	httpServer := http.Server{
		Addr:              *listen,
		Handler:           server.Handler(),
		ReadHeaderTimeout: 5 * time.Second,
	}

	log.Printf("netstable listening on %s", *listen)
	log.Printf("records file: %s", *dataPath)
	if *bandwidthLimitMbps > 0 {
		log.Printf("bandwidth limit: %.2f Mbps", *bandwidthLimitMbps)
	} else {
		log.Printf("bandwidth limit: unlimited")
	}
	if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		log.Fatalf("server failed: %v", err)
	}
}

func runClient(args []string) {
	flags := flag.NewFlagSet("client", flag.ExitOnError)
	serverURL := flags.String("server", "", "NetStable server URL, for example http://103.46.93.38:18080")
	mode := flags.String("mode", "http", "client mode: http or iperf3")
	target := flags.String("target", "", "target label saved in records; defaults to server host")
	durationSeconds := flags.Int("duration", 30, "seconds for each phase: download first, then upload")
	timeoutMillis := flags.Int("timeout", 3000, "request timeout in milliseconds")
	downloadBytes := flags.Int("download-bytes", appclient.DefaultDownloadBytes, "bytes per download request")
	uploadBytes := flags.Int("upload-bytes", appclient.DefaultUploadBytes, "bytes per upload request")
	protocol := flags.String("protocol", "tcp", "iperf3 protocol: tcp or udp")
	reverse := flags.Bool("reverse", false, "iperf3 reverse mode: server sends to client")
	bandwidthMbps := flags.Float64("bandwidth-mbps", 0, "iperf3 target bandwidth in Mbps; defaults to server limit")
	iperf3Path := flags.String("iperf3-path", "iperf3", "iperf3 executable path for iperf3 mode")
	jsonOutput := flags.Bool("json", false, "print summary as JSON")
	quiet := flags.Bool("quiet", false, "suppress sample progress output")
	_ = flags.Parse(args)

	if *serverURL == "" {
		_, _ = fmt.Fprintln(os.Stderr, "client requires -server, for example: netstable client -server http://103.46.93.38:18080")
		flags.Usage()
		os.Exit(2)
	}

	if strings.EqualFold(*mode, "iperf3") {
		runIperf3Client(*serverURL, *target, *durationSeconds, *timeoutMillis, *protocol, *reverse, *bandwidthMbps, *iperf3Path, *jsonOutput, *quiet)
		return
	}
	if !strings.EqualFold(*mode, "http") {
		_, _ = fmt.Fprintf(os.Stderr, "unknown client mode %q; use http or iperf3\n", *mode)
		os.Exit(2)
	}

	options := appclient.Options{
		ServerURL:     *serverURL,
		Target:        *target,
		Duration:      time.Duration(*durationSeconds) * time.Second,
		Timeout:       time.Duration(*timeoutMillis) * time.Millisecond,
		DownloadBytes: *downloadBytes,
		UploadBytes:   *uploadBytes,
		OnQueueWaiting: func(message string) {
			_, _ = fmt.Fprintln(os.Stderr, message)
		},
	}
	if !*quiet && !*jsonOutput {
		options.OnSample = func(sample probe.Sample) {
			printSample(os.Stderr, sample)
		}
	}
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
	defer stop()
	summary, err := appclient.Run(ctx, options)
	if err != nil {
		_, _ = fmt.Fprintf(os.Stderr, "client failed: %v\n", err)
		os.Exit(1)
	}

	if *jsonOutput {
		_ = json.NewEncoder(os.Stdout).Encode(summary)
		return
	}
	printSummary(os.Stdout, summary)
}

func runIperf3Client(serverURL string, target string, durationSeconds int, timeoutMillis int, protocol string, reverse bool, bandwidthMbps float64, iperf3Path string, jsonOutput bool, quiet bool) {
	options := appclient.Iperf3Options{
		ServerURL:     serverURL,
		Target:        target,
		Duration:      time.Duration(durationSeconds) * time.Second,
		Timeout:       time.Duration(timeoutMillis) * time.Millisecond,
		Protocol:      protocol,
		Reverse:       reverse,
		BandwidthMbps: bandwidthMbps,
		Iperf3Path:    iperf3Path,
	}
	if !quiet && !jsonOutput {
		options.OnSample = func(sample probe.Sample) {
			printSample(os.Stderr, sample)
		}
	}
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
	defer stop()
	summary, err := appclient.RunIperf3(ctx, options)
	if err != nil {
		_, _ = fmt.Fprintf(os.Stderr, "iperf3 client failed: %v\n", err)
		os.Exit(1)
	}
	if jsonOutput {
		_ = json.NewEncoder(os.Stdout).Encode(summary)
		return
	}
	printSummary(os.Stdout, summary)
}

func printSample(out *os.File, sample probe.Sample) {
	status := "OK"
	if !sample.Success {
		status = "FAIL"
	}
	switch sample.Kind {
	case "latency":
		_, _ = fmt.Fprintf(out, "%s %s %.1f ms %s\n", sample.Kind, status, sample.LatencyMS, sample.Error)
	case "download", "upload":
		_, _ = fmt.Fprintf(out, "%s %s 瞬时 %.2f Mbps %d bytes %s\n", sample.Kind, status, sample.Mbps, sample.Bytes, sample.Error)
	default:
		_, _ = fmt.Fprintf(out, "%s %s %s\n", sample.Kind, status, sample.Error)
	}
}

func printSummary(out *os.File, summary probe.Summary) {
	_, _ = fmt.Fprintf(out, "目标: %s\n", summary.Target)
	_, _ = fmt.Fprintf(out, "评分: %.2f\n", summary.Score)
	_, _ = fmt.Fprintf(out, "下载: avg %.2f Mbps, max %.2f Mbps\n", summary.AvgDownloadMbps, summary.MaxDownloadMbps)
	_, _ = fmt.Fprintf(out, "上传: avg %.2f Mbps, max %.2f Mbps\n", summary.AvgUploadMbps, summary.MaxUploadMbps)
	_, _ = fmt.Fprintf(out, "延迟: avg %.2f ms, min %.2f ms, max %.2f ms, jitter %.2f ms\n", summary.AvgLatencyMS, summary.MinLatencyMS, summary.MaxLatencyMS, summary.JitterMS)
	_, _ = fmt.Fprintf(out, "样本: %d, 成功: %d, 失败: %d, 丢失率: %.2f%%\n", summary.Samples, summary.Successes, summary.Failures, summary.PacketLossPct)
}

func chainResolvers(resolvers ...appweb.ClientResolver) appweb.ClientResolver {
	return func(ctx context.Context, ip string) (recordstore.ClientInfo, error) {
		var firstErr error
		var best recordstore.ClientInfo
		for _, resolver := range resolvers {
			if resolver == nil {
				continue
			}
			info, err := resolver(ctx, ip)
			if err != nil {
				if firstErr == nil {
					firstErr = err
				}
				continue
			}
			if info.ISP != "" {
				return info, nil
			}
			if best.IP == "" {
				best = info
			}
		}
		if best.IP != "" {
			return best, nil
		}
		return recordstore.ClientInfo{}, firstErr
	}
}