rtu_linux_services/web-ui.service/web_ui.go

package main

import (
	"compress/gzip"
	"context"
	"embed"
	"io/fs"
	"net/http"
	"strings"
	"time"
)

//go:embed web/*
var webRootFS embed.FS

type WebUI struct {
	fs     fs.FS
	mux    *http.ServeMux
	server *http.Server
}

type gzipResponseWriter struct {
	http.ResponseWriter
	writer *gzip.Writer
}

func (w gzipResponseWriter) Write(b []byte) (int, error) {
	return w.writer.Write(b)
}

func gzipHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodHead ||
			!strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
			h.ServeHTTP(w, r)
			return
		}

		path := strings.ToLower(r.URL.Path)

		if !(strings.HasSuffix(path, ".html") || strings.HasSuffix(path, ".js") ||
			strings.HasSuffix(path, ".css")) {
			h.ServeHTTP(w, r)
			return
		}

		gz, err := gzip.NewWriterLevel(w, gzip.DefaultCompression)
		if err != nil {
			h.ServeHTTP(w, r)
			return
		}
		defer gz.Close()

		w.Header().Set("Content-Encoding", "gzip")
		w.Header().Set("Vary", "Accept-Encoding")
		w.Header().Del("Content-Length")

		h.ServeHTTP(&gzipResponseWriter{ResponseWriter: w, writer: gz}, r)
	})
}

func staticHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		path := r.URL.Path
		if strings.HasPrefix(path, "/static/") {
			w.Header().Set("Cache-Control", "public, max-age=86400")
		}
		h.ServeHTTP(w, r)
	})
}

func NewWebUI(addr string) (*WebUI, error) {
	sub, err := fs.Sub(webRootFS, "web")
	if err != nil {
		return nil, err
	}

	mux := http.NewServeMux()

	static := http.FileServer(http.FS(sub))

	mux.Handle(
		"/static/",
		staticHandler(static),
	)

	registerRoutes(mux)

	return &WebUI{
		fs:  sub,
		mux: mux,
		server: &http.Server{
			Addr:    addr,
			Handler: gzipHandler(mux),
		},
	}, nil
}

func (ui *WebUI) Start() error {
	startSessionCleaner() // session过期自动回收
	ui.mux.HandleFunc("/", ui.rootHandler)
	return ui.server.ListenAndServe()
}

func (ui *WebUI) Stop() error {
	if ui.server == nil {
		return nil
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	err := ui.server.Shutdown(ctx)
	if err != nil {
		_ = ui.server.Close()
	}

	return err
}