lodge/backend/internal/server/server.go

package server

import (
	"context"
	"crypto/sha256"
	"crypto/subtle"
	_ "embed"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"lodge/internal/protocol"
)

const (
	heartbeat     = 30 * time.Second
	sessionCookie = "lodge_session"
	loginDelay    = time.Second
)

var (
	//go:embed install.sh
	installScript string

	repoKeyPattern = regexp.MustCompile(`^(tmp-)?[0-9a-f]{40,64}$`)
	hashPattern    = regexp.MustCompile(`^[0-9a-f]{64}$`)
	codePattern    = regexp.MustCompile(`^[0-9a-f]+$`)
)

type Config struct {
	AdminToken string
	PublicURL  string
	StaticDir  string
	ClientsDir string
}

type Server struct {
	config Config
	store  *Store

	mu          sync.Mutex
	subscribers map[chan protocol.Event]struct{}
	streams     map[int64]int
	seen        map[int64]int64
}

type machineKey struct{}

func New(store *Store, config Config) http.Handler {
	s := &Server{config: config, store: store, subscribers: map[chan protocol.Event]struct{}{}, streams: map[int64]int{}, seen: map[int64]int64{}}
	mux := http.NewServeMux()
	mux.HandleFunc("GET /i/{code}", s.install)
	mux.Handle("GET /dl/", http.StripPrefix("/dl/", http.FileServer(http.Dir(config.ClientsDir))))
	mux.HandleFunc("POST /api/claim", s.claim)
	mux.HandleFunc("POST /api/session", s.login)
	mux.HandleFunc("GET /api/session", s.admin(noContent))
	mux.HandleFunc("DELETE /api/session", s.logout)
	mux.HandleFunc("POST /api/codes", s.admin(s.createCode))
	mux.HandleFunc("GET /api/machines", s.admin(s.machines))
	mux.HandleFunc("DELETE /api/machines/{id}", s.admin(s.deleteMachine))
	mux.HandleFunc("DELETE /api/repos/{key}", s.admin(s.deleteRepo))
	mux.HandleFunc("GET /api/repos", s.auth(s.repos))
	mux.HandleFunc("GET /api/repos/{key}/state", s.auth(s.state))
	mux.HandleFunc("GET /api/repos/{key}/revisions", s.auth(s.revisions))
	mux.HandleFunc("POST /api/repos/{key}/revisions", s.auth(s.push))
	mux.HandleFunc("GET /api/blobs/{hash}", s.auth(s.blob))
	mux.HandleFunc("PUT /api/blobs/{hash}", s.auth(s.putBlob))
	mux.HandleFunc("GET /api/events", s.auth(s.events))
	mux.HandleFunc("GET /api/", http.NotFound)
	mux.HandleFunc("GET /", s.static)
	return mux
}

func writeJSON(w http.ResponseWriter, status int, value any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(value)
}

func writeError(w http.ResponseWriter, err error) {
	status := http.StatusInternalServerError
	switch {
	case errors.Is(err, ErrConflict):
		status = http.StatusConflict
	case errors.Is(err, ErrNotFound):
		status = http.StatusNotFound
	}
	http.Error(w, err.Error(), status)
}

func (s *Server) isAdminToken(token string) bool {
	return s.config.AdminToken != "" && subtle.ConstantTimeCompare([]byte(token), []byte(s.config.AdminToken)) == 1
}

func (s *Server) isAdmin(r *http.Request) bool {
	if cookie, err := r.Cookie(sessionCookie); err == nil && s.store.SessionValid(r.Context(), cookie.Value) {
		return true
	}
	return s.isAdminToken(bearer(r))
}

func noContent(w http.ResponseWriter, _ *http.Request) {
	w.WriteHeader(http.StatusNoContent)
}

func bearer(r *http.Request) string {
	token, _ := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer ")
	return token
}

func (s *Server) admin(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if !s.isAdmin(r) {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		next(w, r)
	}
}

func (s *Server) auth(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		var machineID int64
		if !s.isAdmin(r) {
			id, err := s.store.MachineID(r.Context(), bearer(r))
			if err != nil {
				http.Error(w, "unauthorized", http.StatusUnauthorized)
				return
			}
			machineID = id
			s.track(id, 0)
		}
		if key := r.PathValue("key"); key != "" && !repoKeyPattern.MatchString(key) {
			http.Error(w, "bad repository key", http.StatusBadRequest)
			return
		}
		next(w, r.WithContext(context.WithValue(r.Context(), machineKey{}, machineID)))
	}
}

func (s *Server) publicURL(r *http.Request) string {
	if s.config.PublicURL != "" {
		return strings.TrimSuffix(s.config.PublicURL, "/")
	}
	scheme := r.Header.Get("X-Forwarded-Proto")
	if scheme == "" {
		scheme = "http"
	}
	return scheme + "://" + r.Host
}

func (s *Server) install(w http.ResponseWriter, r *http.Request) {
	code := r.PathValue("code")
	if !codePattern.MatchString(code) || !s.store.CodeValid(r.Context(), code) {
		http.Error(w, "echo 'lodge: this install link has expired, create a new one' >&2; exit 1", http.StatusNotFound)
		return
	}
	w.Header().Set("Content-Type", "text/x-shellscript")
	_, _ = strings.NewReplacer("{{URL}}", s.publicURL(r), "{{CODE}}", code).WriteString(w, installScript)
}

func (s *Server) setSession(w http.ResponseWriter, r *http.Request, token string, expires time.Time) {
	http.SetCookie(w, &http.Cookie{
		Name:     sessionCookie,
		Value:    token,
		Path:     "/",
		Expires:  expires,
		HttpOnly: true,
		Secure:   strings.HasPrefix(s.publicURL(r), "https://"),
		SameSite: http.SameSiteStrictMode,
	})
}

func (s *Server) login(w http.ResponseWriter, r *http.Request) {
	var req protocol.Login
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if !s.isAdminToken(req.Password) {
		time.Sleep(loginDelay)
		http.Error(w, "wrong password", http.StatusUnauthorized)
		return
	}
	token, expires, err := s.store.CreateSession(r.Context())
	if err != nil {
		writeError(w, err)
		return
	}
	s.setSession(w, r, token, expires)
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) logout(w http.ResponseWriter, r *http.Request) {
	if cookie, err := r.Cookie(sessionCookie); err == nil {
		_ = s.store.DeleteSession(r.Context(), cookie.Value)
	}
	s.setSession(w, r, "", time.Unix(0, 0))
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) track(machineID int64, streams int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.seen[machineID] = time.Now().UnixMilli()
	s.streams[machineID] += streams
}

func (s *Server) machines(w http.ResponseWriter, r *http.Request) {
	machines, err := s.store.Machines(r.Context())
	if err != nil {
		writeError(w, err)
		return
	}
	s.mu.Lock()
	for i, machine := range machines {
		machines[i].Online = s.streams[machine.ID] > 0
		machines[i].LastSeen = max(s.seen[machine.ID], machine.LastPush)
	}
	s.mu.Unlock()
	writeJSON(w, http.StatusOK, machines)
}

func (s *Server) deleteMachine(w http.ResponseWriter, r *http.Request) {
	id, err := strconv.ParseInt(r.PathValue("id"), 10, 64)
	if err == nil {
		err = s.store.DeleteMachine(r.Context(), id)
	}
	if err != nil {
		writeError(w, ErrNotFound)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) deleteRepo(w http.ResponseWriter, r *http.Request) {
	key := r.PathValue("key")
	if !repoKeyPattern.MatchString(key) {
		http.Error(w, "bad repository key", http.StatusBadRequest)
		return
	}
	if err := s.store.DeleteRepo(r.Context(), key); err != nil {
		writeError(w, err)
		return
	}
	s.publish(protocol.Event{Repo: key})
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) createCode(w http.ResponseWriter, r *http.Request) {
	code, expires, err := s.store.CreateCode(r.Context())
	if err != nil {
		writeError(w, err)
		return
	}
	writeJSON(w, http.StatusCreated, protocol.CodeResponse{
		Code:      code,
		Command:   fmt.Sprintf("curl -fsSL %s/i/%s | sh", s.publicURL(r), code),
		ExpiresAt: expires.UnixMilli(),
	})
}

func (s *Server) claim(w http.ResponseWriter, r *http.Request) {
	var req protocol.ClaimRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	token, err := s.store.Claim(r.Context(), req.Code, req.Name)
	if err != nil {
		writeError(w, err)
		return
	}
	writeJSON(w, http.StatusCreated, protocol.ClaimResponse{Token: token})
}

func (s *Server) repos(w http.ResponseWriter, r *http.Request) {
	repos, err := s.store.Repos(r.Context())
	if err != nil {
		writeError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, repos)
}

func (s *Server) state(w http.ResponseWriter, r *http.Request) {
	state, err := s.store.State(r.Context(), r.PathValue("key"), r.URL.Query().Get("at"))
	if err != nil {
		writeError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, state)
}

func (s *Server) revisions(w http.ResponseWriter, r *http.Request) {
	key, query := r.PathValue("key"), r.URL.Query()
	var revisions []protocol.Revision
	var err error
	if query.Has("after") {
		after, parseErr := strconv.ParseInt(query.Get("after"), 10, 64)
		if parseErr != nil {
			http.Error(w, "bad after", http.StatusBadRequest)
			return
		}
		revisions, err = s.store.RevisionsAfter(r.Context(), key, query.Get("path"), after)
	} else {
		revisions, err = s.store.Revisions(r.Context(), key, query.Get("path"))
	}
	if err != nil {
		writeError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, revisions)
}

func (s *Server) push(w http.ResponseWriter, r *http.Request) {
	var push protocol.Push
	if err := json.NewDecoder(r.Body).Decode(&push); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if len(push.Changes) == 0 {
		http.Error(w, "no changes", http.StatusBadRequest)
		return
	}
	for _, change := range push.Changes {
		validHash := change.Hash == "" || hashPattern.MatchString(change.Hash)
		if !protocol.ValidSpace(change.Space) || !protocol.SafePath(change.Path) || !validHash || change.Hash == change.Base {
			http.Error(w, "bad change for "+change.Path, http.StatusBadRequest)
			return
		}
	}
	key := r.PathValue("key")
	machineID, _ := r.Context().Value(machineKey{}).(int64)
	seq, err := s.store.Push(r.Context(), key, machineID, push)
	if err != nil {
		writeError(w, err)
		return
	}
	event := protocol.Event{Repo: key, Seq: seq, Author: push.Author}
	s.publish(event)
	writeJSON(w, http.StatusCreated, event)
}

func (s *Server) blob(w http.ResponseWriter, r *http.Request) {
	data, err := s.store.Blob(r.Context(), r.PathValue("hash"))
	if err != nil {
		writeError(w, err)
		return
	}
	w.Header().Set("Content-Type", "application/octet-stream")
	w.Header().Set("Cache-Control", "private, max-age=31536000, immutable")
	_, _ = w.Write(data)
}

func (s *Server) putBlob(w http.ResponseWriter, r *http.Request) {
	data, err := io.ReadAll(http.MaxBytesReader(w, r.Body, protocol.MaxBlobSize))
	if err != nil {
		http.Error(w, err.Error(), http.StatusRequestEntityTooLarge)
		return
	}
	sum := sha256.Sum256(data)
	hash := hex.EncodeToString(sum[:])
	if hash != r.PathValue("hash") {
		http.Error(w, "content does not match hash", http.StatusBadRequest)
		return
	}
	if err := s.store.PutBlob(r.Context(), hash, data); err != nil {
		writeError(w, err)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) publish(event protocol.Event) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for subscriber := range s.subscribers {
		select {
		case subscriber <- event:
		default:
		}
	}
}

func (s *Server) events(w http.ResponseWriter, r *http.Request) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "streaming unsupported", http.StatusInternalServerError)
		return
	}
	machineID, _ := r.Context().Value(machineKey{}).(int64)
	subscriber := make(chan protocol.Event, 16)
	s.mu.Lock()
	s.subscribers[subscriber] = struct{}{}
	s.mu.Unlock()
	s.track(machineID, 1)
	defer func() {
		s.mu.Lock()
		delete(s.subscribers, subscriber)
		s.mu.Unlock()
		s.track(machineID, -1)
	}()

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	fmt.Fprint(w, ": connected\n\n")
	flusher.Flush()
	ticker := time.NewTicker(heartbeat)
	defer ticker.Stop()
	for {
		select {
		case <-r.Context().Done():
			return
		case <-ticker.C:
			s.track(machineID, 0)
			fmt.Fprint(w, ": ping\n\n")
		case event := <-subscriber:
			data, _ := json.Marshal(event)
			fmt.Fprintf(w, "data: %s\n\n", data)
		}
		flusher.Flush()
	}
}

func (s *Server) static(w http.ResponseWriter, r *http.Request) {
	file := filepath.Join(s.config.StaticDir, filepath.FromSlash(filepath.Clean("/"+r.URL.Path)))
	if info, err := os.Stat(file); err != nil || info.IsDir() {
		file = filepath.Join(s.config.StaticDir, "index.html")
	}
	http.ServeFile(w, r, file)
}