Commilitia-Drop/desktop/platform/sessioninject.go

package platform

import (
	"bytes"
	"encoding/json"
	"net/http"
	"strconv"

	"github.com/wailsapp/wails/v2/pkg/options/assetserver"
)

// bootPayload is the state injected into the page at load so the WebView store
// can hydrate synchronously: the restored session (refresh-free — see
// SessionView), the device name (both of which the wails:// scheme can't persist
// in WebView storage), and the API base the page should fetch from.
type bootPayload struct {
	Session    *SessionView `json:"session"`
	DeviceName string       `json:"device_name"`
	// APIBase points the page at a loopback HTTP origin for /api on Windows,
	// where the custom-scheme AssetServer buffers streaming responses (so SSE
	// never arrives). Empty on macOS — the page then uses same-origin relative
	// /api through the custom scheme, which streams fine there.
	APIBase string `json:"api_base,omitempty"`
	// DeviceType is the native client kind (macos / windows / linux); the page
	// forwards it as X-Device-Type so the backend lists it correctly instead of
	// defaulting every device to "browser".
	DeviceType string `json:"device_type,omitempty"`
}

// SessionInjectMiddleware injects the boot state into the served index.html as
// `window.__CDROP_BOOT__`, so the WebView store hydrates synchronously at startup
// — no WebView storage (which doesn't survive restart for the wails:// scheme)
// and no binding round-trip (window.go isn't ready that early). Read once at app
// launch. device_name + api_base are injected always (needed even logged out);
// the session is included only when one was restored.
//
// Only the top-level HTML document is rewritten; assets and the /api proxy
// (including the SSE stream) pass straight through untouched. The JSON is
// produced by encoding/json with default HTML escaping, so a display name
// containing "</script>" can't break out of the inline script.
func SessionInjectMiddleware(session *LoginResult, deviceName, apiBase, deviceType string) assetserver.Middleware {
	boot := bootPayload{DeviceName: deviceName, APIBase: apiBase, DeviceType: deviceType}
	if session != nil {
		view := session.View() // strip the refresh_token before it ever reaches the page
		boot.Session = &view
	}
	payload, _ := json.Marshal(boot)

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.Method != http.MethodGet || !isDocumentRequest(r.URL.Path) {
				next.ServeHTTP(w, r)
				return
			}

			// Force a full 200 (never a 304): strip conditional headers so
			// http.ServeContent always returns the whole document for us to
			// inject. Without this, a WebView that cached the first (logged-out)
			// page could send If-Modified-Since on the next launch, get a 304 with
			// an empty body, reuse its stale page, and the session wouldn't restore.
			r.Header.Del("If-Modified-Since")
			r.Header.Del("If-None-Match")
			r.Header.Del("If-Range")

			rec := &captureWriter{header: http.Header{}}
			next.ServeHTTP(rec, r)
			body := rec.buf.Bytes()

			if idx := bytes.Index(body, []byte("</head>")); idx != -1 {
				script := []byte("<script>window.__CDROP_BOOT__=" + string(payload) + ";</script>")
				out := make([]byte, 0, len(body)+len(script))
				out = append(out, body[:idx]...)
				out = append(out, script...)
				out = append(out, body[idx:]...)
				body = out
			}

			h := w.Header()
			for k, vs := range rec.header {
				for _, v := range vs {
					h.Add(k, v)
				}
			}
			// Never let the WebView cache the document — the injected session
			// changes between launches, so a cached copy must not be reused.
			h.Del("Last-Modified")
			h.Del("ETag")
			h.Set("Cache-Control", "no-store, no-cache, must-revalidate")
			h.Set("Pragma", "no-cache")
			h.Set("Content-Length", strconv.Itoa(len(body)))
			code := rec.code
			if code == 0 {
				code = http.StatusOK
			}
			w.WriteHeader(code)
			_, _ = w.Write(body)
		})
	}
}

// isDocumentRequest matches the top-level HTML document the WebView loads.
func isDocumentRequest(path string) bool {
	return path == "" || path == "/" || path == "/index.html"
}

// captureWriter buffers a handler's response so the middleware can rewrite the
// HTML body before forwarding it. Only used for the small index.html document.
type captureWriter struct {
	header http.Header
	buf    bytes.Buffer
	code   int
}

func (c *captureWriter) Header() http.Header        { return c.header }
func (c *captureWriter) WriteHeader(code int)       { c.code = code }
func (c *captureWriter) Write(b []byte) (int, error) { return c.buf.Write(b) }