web-vnc/internal/server/server.go

// Package server wires HTTP/WS routes, session middleware and static assets.
package server

import (
	"embed"
	"encoding/json"
	"io/fs"
	"log"
	"net"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/web-vnc/internal/auth"
	"github.com/web-vnc/internal/config"
	"github.com/web-vnc/internal/relay"
)

//go:embed all:static
var staticFS embed.FS

// Server is the configured HTTP server for the web-vnc gateway.
type Server struct {
	cfg     config.Config
	auth    *auth.Service
	relay   *relay.Server
	static  fs.FS
	cookies *htmlTemplate
}

// New builds a Server from configuration.
func New(cfg config.Config, authSvc *auth.Service, relaySrv *relay.Server) (*Server, error) {
	sub, err := fs.Sub(staticFS, "static")
	if err != nil {
		return nil, err
	}
	if cfg.WebRoot != "" {
		sub = os.DirFS(cfg.WebRoot)
	}
	s := &Server{
		cfg:     cfg,
		auth:    authSvc,
		relay:   relaySrv,
		static:  sub,
		cookies: mustParseLoginTemplate(),
	}
	return s, nil
}

// Handler returns the configured http.Handler.
func (s *Server) Handler() http.Handler {
	mux := http.NewServeMux()

	mux.HandleFunc(s.cfg.LoginPath, s.handleLogin)
	mux.HandleFunc(s.cfg.LogoutPath, s.handleLogout)
	mux.Handle(s.cfg.StatusPath, s.requireSession(http.HandlerFunc(s.handleStatus)))

	mux.Handle(s.cfg.RelayPath, s.requireSession(http.HandlerFunc(s.handleRelay)))
	mux.Handle(s.cfg.NoVNCPath, s.requireSession(http.HandlerFunc(s.serveNoVNCPage)))

	// favicon is fetched by the browser on every page (incl. the login form);
	// serve it without a session so it doesn't redirect to /login and 404.
	mux.HandleFunc("/favicon.ico", s.handleFavicon)

	mux.Handle("/", s.requireSession(http.HandlerFunc(s.handleIndexOrStatic)))

	return s.logRequest(mux)
}

func (s *Server) handleLogin(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodGet {
		s.renderLogin(w, "")
		return
	}
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	ip := clientIP(r)
	if !s.auth.AllowLogin(ip) {
		w.WriteHeader(http.StatusTooManyRequests)
		s.renderLogin(w, "Too many attempts. Try again later.")
		return
	}
	password := r.PostFormValue("password")
	if !s.auth.CheckPassword(password) {
		w.WriteHeader(http.StatusUnauthorized)
		s.renderLogin(w, "Wrong password")
		return
	}
	token := s.auth.IssueSession(time.Now())
	s.auth.SetSessionCookie(w, token)
	http.Redirect(w, r, s.cfg.NoVNCPath, http.StatusSeeOther)
}

func (s *Server) handleLogout(w http.ResponseWriter, r *http.Request) {
	s.auth.ClearSessionCookie(w)
	http.Redirect(w, r, s.cfg.LoginPath, http.StatusSeeOther)
}

// handleStatus reports whether the upstream VNC server is reachable.
func (s *Server) handleStatus(w http.ResponseWriter, r *http.Request) {
	ok := pingVNC(s.cfg.VNCAddr, 1500*time.Millisecond)
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(map[string]any{
		"vnc":         ok,
		"vncAddr":     s.cfg.VNCAddr,
		"vncPassword": s.cfg.VNCPassword,
		"spawned":     s.cfg.SpawnVNC || s.cfg.SpawnCommand != "",
	})
}

func (s *Server) handleRelay(w http.ResponseWriter, r *http.Request) {
	s.relay.ServeHTTP(w, r)
}

// handleFavicon serves the embedded favicon.ico with a long cache. It is
// registered outside requireSession so it loads on the login page too.
func (s *Server) handleFavicon(w http.ResponseWriter, r *http.Request) {
	data, err := fs.ReadFile(s.static, "favicon.ico")
	if err != nil {
		http.NotFound(w, r)
		return
	}
	w.Header().Set("Content-Type", "image/x-icon")
	w.Header().Set("Cache-Control", "public, max-age=604800")
	_, _ = w.Write(data)
}

func (s *Server) serveNoVNCPage(w http.ResponseWriter, r *http.Request) {
	if data, err := fs.ReadFile(s.static, "vnc.html"); err == nil {
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		_, _ = w.Write(data)
		return
	}
	http.NotFound(w, r)
}

func (s *Server) handleIndexOrStatic(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/" {
		http.Redirect(w, r, s.cfg.NoVNCPath, http.StatusSeeOther)
		return
	}
	clean := strings.TrimPrefix(r.URL.Path, "/")
	if clean != "" && !strings.HasPrefix(clean, ".") {
		http.ServeFileFS(w, r, s.static, clean)
		return
	}
	http.NotFound(w, r)
}

func (s *Server) requireSession(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !s.auth.SessionFromRequest(r, time.Now()) {
			if isWebSocket(r) {
				http.Error(w, "unauthorized", http.StatusUnauthorized)
				return
			}
			http.Redirect(w, r, s.cfg.LoginPath, http.StatusSeeOther)
			return
		}
		h.ServeHTTP(w, r)
	})
}

// pingVNC reports whether a TCP connection to addr succeeds within the timeout.
func pingVNC(addr string, timeout time.Duration) bool {
	conn, err := net.DialTimeout("tcp", addr, timeout)
	if err != nil {
		return false
	}
	_ = conn.Close()
	return true
}

func isWebSocket(r *http.Request) bool {
	return strings.EqualFold(r.Header.Get("Upgrade"), "websocket")
}

func clientIP(r *http.Request) string {
	if h := r.Header.Get("X-Forwarded-For"); h != "" {
		return strings.TrimSpace(strings.Split(h, ",")[0])
	}
	host := r.RemoteAddr
	if i := strings.LastIndex(host, ":"); i > 0 {
		host = host[:i]
	}
	return host
}

func (s *Server) logRequest(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		h.ServeHTTP(w, r)
		log.Printf("%s %s %s", r.Method, r.URL.Path, time.Since(start))
	})
}