SimpleRemoter/server/go/web/server.go

package web

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"html"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
	"unicode"
	"unicode/utf8"

	"github.com/yuanyuanxiang/SimpleRemoter/server/go/hub"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/logger"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/wsauth"
)

// Server serves the web remote desktop UI: the embedded index.html, xterm.js
// static assets, the PWA manifest, and JSON APIs backed by the device hub.
// WebSocket signaling and screen streaming will be wired up in later phases.
type Server struct {
	port              int
	log               *logger.Logger
	srv               *http.Server
	hub               *hub.Hub
	auth              *wsauth.Authenticator
	ws                *wsHub
	allowedOrigins    []string // for WS Origin allowlist; empty = same-origin only
	loginIPLimit      *wsauth.RateLimiter
	loginUserLimit    *wsauth.RateLimiter
	trustForwardedFor bool // honor X-Forwarded-For (behind trusted proxy only)
	basePath          string // web UI mount prefix (e.g. "/remoter"); "" = domain root
	brandName         string // white-label title; "" = "SimpleRemoter" (official branding)
	sso               SSODetector // nil = SSO disabled
	sweepDone         chan struct{} // closes to stop the login-limiter sweeper
	sweepStopOnce     sync.Once
}

// Config tunes the server's exposed-on-public-HTTPS hardening knobs.
// All fields are optional; zero values pick reasonable defaults.
type Config struct {
	// AllowedOrigins is the comma-separated list of Origin header values
	// the WebSocket upgrade will accept in addition to same-origin
	// requests. Empty (default) → only same-origin upgrades are allowed,
	// which is correct when the web UI and the WS endpoint are served
	// from the same host.
	AllowedOrigins []string
	// LoginIPLimit / LoginUserLimit throttle the get_salt + login flow
	// per source IP and per username respectively. Pass nil to disable
	// either dimension (e.g. dev mode).
	LoginIPLimit   *wsauth.RateLimiter
	LoginUserLimit *wsauth.RateLimiter
	// TrustForwardedFor switches client-IP extraction from RemoteAddr
	// (default) to the last entry of X-Forwarded-For. Set true only when
	// running behind a reverse proxy that you control; on direct
	// exposure the header is client-controlled and would let attackers
	// evade per-IP rate limits.
	TrustForwardedFor bool
	// BasePath mounts the whole web UI (index, manifest, static assets,
	// /health, /api/devices, /ws) under a URL prefix, so a reverse proxy
	// can serve it from a subpath of an existing domain (e.g.
	// https://example.com/remoter). Must be empty or a leading-slash,
	// no-trailing-slash path; "" (default) serves from the domain root.
	BasePath string
	// BrandName is the white-label product title shown in index.html and the
	// PWA manifest when BasePath is non-empty (i.e. mounted under another
	// domain). Read from the YAMA_WEB_BRAND_NAME env var by cmd. Empty leaves
	// the official "SimpleRemoter" branding in place.
	BrandName string
	// SSO is an optional detector that maps an already-authenticated portal
	// session (cookie/header) into a YAMA web session via /sso. nil disables
	// SSO and the front-end falls back to the built-in login form.
	SSO SSODetector
}

// New creates an HTTP server bound to the given port. port=0 disables the server.
// The hub provides read access to the online-device registry; the authenticator
// owns user accounts and session tokens.
func New(port int, log *logger.Logger, h *hub.Hub, auth *wsauth.Authenticator) *Server {
	return &Server{port: port, log: log, hub: h, auth: auth}
}

// WithConfig applies hardening configuration. Returns the receiver for
// chainable setup. Safe to call before Start; ignored thereafter.
func (s *Server) WithConfig(cfg Config) *Server {
	s.allowedOrigins = cfg.AllowedOrigins
	s.loginIPLimit = cfg.LoginIPLimit
	s.loginUserLimit = cfg.LoginUserLimit
	s.trustForwardedFor = cfg.TrustForwardedFor
	s.basePath = cfg.BasePath
	s.brandName = resolveBrandName(cfg.BrandName, s.basePath)
	s.sso = cfg.SSO
	return s
}

// Start launches the server in a goroutine and returns immediately.
// If port is 0, returns nil without starting anything.
func (s *Server) Start() error {
	if s.port == 0 {
		s.log.Info("HTTP server disabled (-http-port=0)")
		return nil
	}

	s.ws = newWSHub(s.auth, s.hub, s.log).
		withOriginAllowlist(s.allowedOrigins).
		withLoginRateLimiters(s.loginIPLimit, s.loginUserLimit).
		withTrustForwardedFor(s.trustForwardedFor)

	mux := http.NewServeMux()
	mux.HandleFunc("/", s.handleIndex)
	mux.HandleFunc("/health", s.handleHealth)
	mux.HandleFunc("/manifest.json", s.handleManifest)
	mux.HandleFunc("/api/devices", s.requireBearer(s.handleDevices))
	mux.HandleFunc("/sso", s.handleSSO)
	mux.HandleFunc("/ws", s.ws.serve)
	mux.HandleFunc("/static/xterm.js", staticHandler(xtermJS, "application/javascript; charset=utf-8"))
	mux.HandleFunc("/static/xterm.css", staticHandler(xtermCSS, "text/css; charset=utf-8"))
	mux.HandleFunc("/static/xterm-fit.js", staticHandler(xtermFitJS, "application/javascript; charset=utf-8"))
	mux.HandleFunc("/static/opus-decoder.min.js", staticHandler(opusDecoderJS, "application/javascript; charset=utf-8"))

	s.srv = &http.Server{
		Addr:              ":" + strconv.Itoa(s.port),
		Handler:           s.mount(mux),
		ReadHeaderTimeout: 10 * time.Second,
	}

	// Bind synchronously so port-in-use / permission errors propagate to the
	// caller instead of being lost inside the goroutine after a misleading
	// "started" log line.
	ln, err := net.Listen("tcp", s.srv.Addr)
	if err != nil {
		return fmt.Errorf("listen on :%d: %w", s.port, err)
	}
	s.log.Info("HTTP server started on :%d", s.port)
	go func() {
		if err := s.srv.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) {
			s.log.Error("HTTP server stopped: %v", err)
		}
	}()

	// Prune aged-out login rate-limit entries so the key maps (especially the
	// username-keyed one, which is attacker-controlled) stay bounded.
	s.sweepDone = make(chan struct{})
	go s.sweepLoginLimiters(s.sweepDone)

	return nil
}

// sweepLoginLimiters periodically sweeps both web login rate limiters. Runs
// until `stop` is closed (from Stop).
func (s *Server) sweepLoginLimiters(stop <-chan struct{}) {
	t := time.NewTicker(time.Minute)
	defer t.Stop()
	for {
		select {
		case <-stop:
			return
		case <-t.C:
			if s.loginIPLimit != nil {
				s.loginIPLimit.Sweep()
			}
			if s.loginUserLimit != nil {
				s.loginUserLimit.Sweep()
			}
		}
	}
}

// Stop gracefully shuts the server down.
func (s *Server) Stop() {
	if s.sweepDone != nil {
		s.sweepStopOnce.Do(func() { close(s.sweepDone) })
	}
	if s.ws != nil {
		s.ws.stop()
	}
	if s.srv == nil {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_ = s.srv.Shutdown(ctx)
}

// basePathPlaceholder is the token in index.html substituted with the
// configured base path at serve time. Empty base path → "" → root URLs.
const basePathPlaceholder = "__BASE_PATH__"

// brandPlaceholder is the token in index.html substituted with the product
// title at serve time (official or white-label, depending on BasePath).
const brandPlaceholder = "__BRAND_NAME__"

// mount wraps the root-path mux with the configured base path. When basePath
// is empty the mux is returned unchanged (domain-root serving, 方案A). Otherwise
// the bare prefix redirects to prefix+"/", paths outside the prefix 404, and
// matching paths are stripped before reaching the mux (方案B, subpath serving).
func (s *Server) mount(h http.Handler) http.Handler {
	if s.basePath == "" {
		return h
	}
	base := s.basePath // normalized by cmd: leading "/", no trailing "/"
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == base {
			http.Redirect(w, r, base+"/", http.StatusMovedPermanently)
			return
		}
		if !strings.HasPrefix(r.URL.Path, base+"/") {
			http.NotFound(w, r)
			return
		}
		http.StripPrefix(base, h).ServeHTTP(w, r)
	})
}

func (s *Server) handleIndex(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" && r.URL.Path != "/index.html" {
		http.NotFound(w, r)
		return
	}
	body := bytes.ReplaceAll(IndexHTML, []byte(basePathPlaceholder), []byte(s.basePath))
	body = bytes.ReplaceAll(body, []byte(brandPlaceholder), []byte(html.EscapeString(s.displayName())))
	if s.isWhiteLabel() {
		body = stripBrandBlock(body, "NOTICE")
		body = stripBrandBlock(body, "LINKS")
	}
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.Header().Set("Cache-Control", "no-cache")
	_, _ = w.Write(body)
}

func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	_, _ = w.Write([]byte(`{"status":"ok"}`))
}

// handleDevices returns a JSON snapshot of currently-online devices. Empty
// array (not null) when no clients are connected — matches what the front-end
// will eventually expect. Auth-gated via requireBearer.
func (s *Server) handleDevices(w http.ResponseWriter, r *http.Request) {
	devices := s.hub.ListDevices()
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Header().Set("Cache-Control", "no-store")
	if err := json.NewEncoder(w).Encode(devices); err != nil {
		s.log.Error("encode /api/devices: %v", err)
	}
}

// requireBearer wraps a handler with `Authorization: Bearer <token>` auth
// against the same session-token store the WebSocket uses. Returns 401 on
// missing / invalid / expired tokens. Used to gate REST endpoints that
// previously fell through with no auth (notably /api/devices, which
// otherwise leaks the full online-device list to anyone on the internet).
func (s *Server) requireBearer(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		const prefix = "Bearer "
		hdr := r.Header.Get("Authorization")
		if !strings.HasPrefix(hdr, prefix) {
			s.unauthorized(w)
			return
		}
		token := strings.TrimSpace(hdr[len(prefix):])
		if token == "" {
			s.unauthorized(w)
			return
		}
		if _, err := s.auth.ValidateToken(token); err != nil {
			s.unauthorized(w)
			return
		}
		next(w, r)
	}
}

func (s *Server) unauthorized(w http.ResponseWriter) {
	w.Header().Set("WWW-Authenticate", `Bearer realm="yama"`)
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusUnauthorized)
	_, _ = w.Write([]byte(`{"error":"unauthorized"}`))
}

// PWA manifest. Referenced by <link rel="manifest"> in index.html.
// start_url is rooted at the configured base path so an installed PWA
// launches into the right location; "/" when served from the domain root.
func (s *Server) manifestJSON() []byte {
	start := "/"
	if s.basePath != "" {
		start = s.basePath + "/"
	}
	name := s.displayName()
	shortName := "Remoter"
	if s.isWhiteLabel() {
		shortName = name
	}
	// Marshal via encoding/json so an operator-supplied brand name containing
	// quotes or backslashes can't produce invalid JSON.
	manifest := struct {
		Name            string `json:"name"`
		ShortName       string `json:"short_name"`
		StartURL        string `json:"start_url"`
		Display         string `json:"display"`
		Orientation     string `json:"orientation"`
		BackgroundColor string `json:"background_color"`
		ThemeColor      string `json:"theme_color"`
	}{
		Name:            name,
		ShortName:       shortName,
		StartURL:        start,
		Display:         "standalone",
		Orientation:     "any",
		BackgroundColor: "#1a1a2e",
		ThemeColor:      "#1a1a2e",
	}
	b, err := json.MarshalIndent(manifest, "", "  ")
	if err != nil {
		// A struct of plain strings cannot fail to marshal.
		return []byte(`{"start_url":"/","display":"standalone"}`)
	}
	return b
}

func (s *Server) handleManifest(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/manifest+json")
	w.Header().Set("Cache-Control", "public, max-age=86400")
	_, _ = w.Write(s.manifestJSON())
}

// displayName returns the product title shown in the UI: the configured brand
// name when set, otherwise the official "SimpleRemoter".
func (s *Server) displayName() string {
	if s.brandName != "" {
		return s.brandName
	}
	return "SimpleRemoter"
}

// isWhiteLabel reports whether the web UI is mounted under another domain's
// subpath (BasePath non-empty). In that case the official SimpleRemoter
// branding (security notice, website links) is stripped and the configured
// brand name is shown instead.
func (s *Server) isWhiteLabel() bool {
	return s.basePath != ""
}

// resolveBrandName picks the effective white-label title: the configured value
// when non-empty, else a fallback derived from the base path's last segment
// (e.g. "/remoter" → "Remoter"). Empty base path means the UI is served at the
// domain root, so it returns "" (official "SimpleRemoter") and ignores any
// configured name — white-label branding only applies to subpath deployment.
func resolveBrandName(configured, basePath string) string {
	if basePath == "" {
		return ""
	}
	if configured != "" {
		return configured
	}
	seg := strings.Trim(basePath, "/")
	if i := strings.LastIndex(seg, "/"); i >= 0 {
		seg = seg[i+1:]
	}
	return capitalizeFirst(seg)
}

// capitalizeFirst upper-cases the first rune, UTF-8 safe (no byte slicing).
func capitalizeFirst(s string) string {
	if s == "" {
		return s
	}
	r, size := utf8.DecodeRuneInString(s)
	return string(unicode.ToUpper(r)) + s[size:]
}

// stripBrandBlock removes the index.html block delimited by
// <!--BRAND-<name>-START--> … <!--BRAND-<name>-END--> (markers included).
func stripBrandBlock(body []byte, name string) []byte {
	start := []byte("<!--BRAND-" + name + "-START-->")
	end := []byte("<!--BRAND-" + name + "-END-->")
	i := bytes.Index(body, start)
	if i < 0 {
		return body
	}
	j := bytes.Index(body[i:], end)
	if j < 0 {
		return body
	}
	j += i + len(end)
	return append(body[:i], body[j:]...)
}

// staticHandler returns an http.HandlerFunc that serves a fixed byte slice
// with the given content-type. Used for embedded third-party assets (xterm.js etc.)
// that change infrequently — 1-day browser cache is fine.
func staticHandler(body []byte, contentType string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", contentType)
		w.Header().Set("Cache-Control", "public, max-age=86400")
		_, _ = w.Write(body)
	}
}