probo/pkg/server/statichandler/statichandler.go

// Copyright (c) 2025-2026 Probo Inc <hello@probo.com>.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

// Package statichandler provides functionality for serving SPA (Single Page Application) frontends.
package statichandler

import (
	"bytes"
	"compress/gzip"
	"crypto/md5"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"net/http"
	"strings"
)

type (
	// FileRenderer renders dynamic content for a given file path. When
	// registered via WithFileRenderer, the server calls it instead of serving
	// the static embedded bytes. The renderer writes the response body to w.
	FileRenderer func(w io.Writer, r *http.Request) error

	Option func(*Server)

	GzipOptions struct {
		EnableFileTypeCheck bool
		FileTypes           []string
	}

	Server struct {
		spaFS         http.FileSystem
		etags         map[string]string
		indexETag     string
		indexContent  []byte
		gzipOptions   GzipOptions
		fileRenderers map[string]FileRenderer
	}
)

// WithFileRenderer registers a dynamic renderer for the given path (e.g.
// "/index.html"). When the server would serve that file, it calls the
// renderer instead. ETag-based caching is disabled for rendered files.
func WithFileRenderer(path string, renderer FileRenderer) Option {
	return func(s *Server) {
		s.fileRenderers[path] = renderer
	}
}

func NewServer(staticFiles fs.FS, distPath string, gzipOptions GzipOptions, opts ...Option) (*Server, error) {
	subFS, err := fs.Sub(staticFiles, distPath)
	if err != nil {
		return nil, err
	}

	etags := make(map[string]string)

	err = fs.WalkDir(
		subFS,
		".",
		func(path string, d fs.DirEntry, err error) error {
			if err != nil {
				return err
			}

			if d.IsDir() {
				return nil
			}

			info, err := d.Info()
			if err != nil {
				return err
			}

			content := make([]byte, info.Size())

			file, err := subFS.Open(path)
			if err != nil {
				return err
			}

			defer func() { _ = file.Close() }()

			_, err = file.Read(content)
			if err != nil {
				return err
			}

			hash := md5.Sum(content)
			etag := hex.EncodeToString(hash[:])
			etags["/"+path] = etag

			return nil
		},
	)
	if err != nil {
		return nil, fmt.Errorf("cannot generate etags: %w", err)
	}

	indexETag, ok := etags["/index.html"]
	if !ok {
		return nil, errors.New("index.html not found")
	}

	indexFile, err := subFS.Open("index.html")
	if err != nil {
		return nil, err
	}

	indexContent, err := io.ReadAll(indexFile)
	if err != nil {
		return nil, err
	}

	s := &Server{
		spaFS:         http.FS(subFS),
		indexETag:     indexETag,
		indexContent:  indexContent,
		etags:         etags,
		gzipOptions:   gzipOptions,
		fileRenderers: make(map[string]FileRenderer),
	}

	for _, opt := range opts {
		opt(s)
	}

	return s, nil
}

func (s *Server) serveIndex(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")

	if renderer, ok := s.fileRenderers["/index.html"]; ok {
		var buf bytes.Buffer
		if err := renderer(&buf, r); err != nil {
			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
			return
		}

		w.WriteHeader(http.StatusOK)
		_, _ = w.Write(buf.Bytes())

		return
	}

	w.Header().Set("ETag", `"`+s.indexETag+`"`)

	if r.Header.Get("If-None-Match") == `"`+s.indexETag+`"` {
		w.WriteHeader(http.StatusNotModified)
		return
	}

	w.WriteHeader(http.StatusOK)
	_, _ = w.Write(s.indexContent)
}

func (s *Server) ServeSPA(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path

	f, err := s.spaFS.Open(path)
	if err != nil {
		s.serveIndex(w, r)
		return
	}

	defer func() { _ = f.Close() }()

	info, err := f.Stat()
	if err != nil {
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}

	if info.IsDir() {
		s.serveIndex(w, r)
		return
	}

	if renderer, ok := s.fileRenderers[path]; ok {
		var buf bytes.Buffer
		if err := renderer(&buf, r); err != nil {
			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write(buf.Bytes())

		return
	}

	etag, ok := s.etags[path]
	if !ok {
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}

	quotedETag := `"` + etag + `"`
	w.Header().Set("ETag", quotedETag)

	if strings.HasSuffix(path, ".js") || strings.HasSuffix(path, ".css") ||
		strings.HasSuffix(path, ".png") || strings.HasSuffix(path, ".jpg") ||
		strings.HasSuffix(path, ".svg") || strings.HasSuffix(path, ".woff2") {
		w.Header().Set("Cache-Control", "public, max-age=31536000, immutable")
	} else {
		w.Header().Set("Cache-Control", "public, max-age=3600")
	}

	if matchETag := r.Header.Get("If-None-Match"); matchETag != "" {
		if matchETag == quotedETag {
			w.WriteHeader(http.StatusNotModified)
			return
		}
	}

	http.FileServer(s.spaFS).ServeHTTP(w, r)
}

type gzipResponseWriter struct {
	io.Writer
	http.ResponseWriter
}

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

func (s *Server) shouldCompressWithGzip(r *http.Request) bool {
	if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
		return false
	}

	if !s.gzipOptions.EnableFileTypeCheck {
		return true
	}

	path := r.URL.Path
	for _, fileType := range s.gzipOptions.FileTypes {
		if strings.HasSuffix(path, fileType) {
			return true
		}
	}

	return false
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if s.shouldCompressWithGzip(r) {
		w.Header().Set("Content-Encoding", "gzip")
		gz := gzip.NewWriter(w)

		defer func() { _ = gz.Close() }()

		gzw := gzipResponseWriter{Writer: gz, ResponseWriter: w}
		s.ServeSPA(gzw, r)

		return
	}

	s.ServeSPA(w, r)
}