prompt/Server/main.go

// Prompt sync server. No authentication: a sync group is identified by an
// 8-digit code (1234-5678) that acts as the shared secret. Whoever has the
// code can read and overwrite the group's data.
//
// Endpoints:
//
//	POST /api/codes           -> {"code":"1234-5678"}        (create group)
//	GET  /api/sync/{code}     -> {"version":N,"data":{...}}  (404 if unknown)
//	GET  /api/sync/{code}/version -> {"version":N}           (cheap poll)
//	PUT  /api/sync/{code}     -> {"version":N+1}             (upsert payload)
//	GET  /healthz             -> 200
//
// Config:
//
//	PROMPT_SYNC_DATABASE_URL  (default postgres://localhost:5432/prompt_sync?sslmode=disable)
//	PROMPT_SYNC_ADDR          (default :8080)
package main

import (
	"context"
	"crypto/rand"
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"strings"
	"syscall"
	"time"

	_ "github.com/lib/pq"
)

const schema = `
CREATE TABLE IF NOT EXISTS sync_groups (
    code       TEXT PRIMARY KEY,
    version    BIGINT NOT NULL DEFAULT 0,
    data       JSONB,
    updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
)`

var codePattern = regexp.MustCompile(`^\d{4}-\d{4}$`)

type Server struct {
	db *sql.DB
}

func main() {
	// -healthcheck pings the local /healthz and exits non-zero on failure. It
	// exists so the container image can run on a shell-less base (distroless),
	// where there is no wget or curl for a HEALTHCHECK to use.
	if len(os.Args) > 1 && os.Args[1] == "-healthcheck" {
		os.Exit(healthcheck())
	}

	dsn := os.Getenv("PROMPT_SYNC_DATABASE_URL")
	if dsn == "" {
		dsn = "postgres://localhost:5432/prompt_sync?sslmode=disable"
	}
	addr := os.Getenv("PROMPT_SYNC_ADDR")
	if addr == "" {
		addr = ":8080"
	}

	db, err := sql.Open("postgres", dsn)
	if err != nil {
		log.Fatalf("open database: %v", err)
	}
	defer db.Close()
	if err := db.Ping(); err != nil {
		log.Fatalf("connect to database: %v", err)
	}
	if _, err := db.Exec(schema); err != nil {
		log.Fatalf("ensure schema: %v", err)
	}

	s := &Server{db: db}
	mux := http.NewServeMux()
	mux.HandleFunc("POST /api/codes", s.createCode)
	mux.HandleFunc("GET /api/sync/{code}", s.getSync)
	mux.HandleFunc("GET /api/sync/{code}/version", s.getVersion)
	mux.HandleFunc("PUT /api/sync/{code}", s.putSync)
	mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
	})

	server := &http.Server{
		Addr:              addr,
		Handler:           withLogging(mux),
		ReadHeaderTimeout: 10 * time.Second,
	}

	// Shut down cleanly so a redeploy (systemd restart, container stop) never
	// cuts off an in-flight push and leaves a client thinking it failed.
	stop := make(chan os.Signal, 1)
	signal.Notify(stop, os.Interrupt, syscall.SIGTERM)
	go func() {
		<-stop
		log.Println("shutting down")
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		_ = server.Shutdown(ctx)
	}()

	log.Printf("prompt-sync listening on %s", addr)
	if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		log.Fatalf("serve: %v", err)
	}
}

// healthcheck returns 0 when the local /healthz answers 200.
func healthcheck() int {
	addr := os.Getenv("PROMPT_SYNC_ADDR")
	if addr == "" {
		addr = ":8080"
	}
	if strings.HasPrefix(addr, ":") {
		addr = "127.0.0.1" + addr
	}
	client := &http.Client{Timeout: 3 * time.Second}
	response, err := client.Get("http://" + addr + "/healthz")
	if err != nil {
		log.Printf("healthcheck: %v", err)
		return 1
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		log.Printf("healthcheck: unexpected status %d", response.StatusCode)
		return 1
	}
	return 0
}

// withLogging keeps a one-line trail of method, path, status and duration so a
// self-hosted deployment can be debugged from journalctl or docker logs.
func withLogging(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		started := time.Now()
		recorder := &statusRecorder{ResponseWriter: w, status: http.StatusOK}
		next.ServeHTTP(recorder, r)
		log.Printf("%s %s %d %s", r.Method, r.URL.Path, recorder.status, time.Since(started).Round(time.Millisecond))
	})
}

type statusRecorder struct {
	http.ResponseWriter
	status int
}

func (r *statusRecorder) WriteHeader(status int) {
	r.status = status
	r.ResponseWriter.WriteHeader(status)
}

// createCode generates a unique 8-digit code, formats it as 1234-5678, and
// inserts an empty group so the code exists before the first PUT.
func (s *Server) createCode(w http.ResponseWriter, _ *http.Request) {
	for range 10 {
		code, err := newCode()
		if err != nil {
			writeError(w, http.StatusInternalServerError, "could not generate code")
			return
		}
		var exists bool
		err = s.db.QueryRow(`SELECT EXISTS(SELECT 1 FROM sync_groups WHERE code = $1)`, code).Scan(&exists)
		if err != nil {
			writeError(w, http.StatusInternalServerError, "database error")
			return
		}
		if exists {
			continue
		}
		if _, err := s.db.Exec(`INSERT INTO sync_groups (code) VALUES ($1)`, code); err != nil {
			writeError(w, http.StatusInternalServerError, "database error")
			return
		}
		writeJSON(w, http.StatusOK, map[string]string{"code": code})
		return
	}
	writeError(w, http.StatusInternalServerError, "could not find a free code")
}

func (s *Server) getSync(w http.ResponseWriter, r *http.Request) {
	code := r.PathValue("code")
	if !codePattern.MatchString(code) {
		writeError(w, http.StatusBadRequest, "code must look like 1234-5678")
		return
	}
	var (
		version   int64
		data      []byte
		updatedAt time.Time
	)
	err := s.db.QueryRow(
		`SELECT version, data, updated_at FROM sync_groups WHERE code = $1`, code,
	).Scan(&version, &data, &updatedAt)
	if err == sql.ErrNoRows {
		writeError(w, http.StatusNotFound, "unknown code")
		return
	}
	if err != nil {
		writeError(w, http.StatusInternalServerError, "database error")
		return
	}
	if data == nil {
		data = []byte("null")
	}
	w.Header().Set("Content-Type", "application/json")
	fmt.Fprintf(w, `{"version":%d,"updatedAt":%q,"data":%s}`,
		version, updatedAt.UTC().Format(time.RFC3339), data)
}

// getVersion is the cheap poll: clients ask for the version every few seconds
// and only download the payload (which can carry base64 images) when it moved.
func (s *Server) getVersion(w http.ResponseWriter, r *http.Request) {
	code := r.PathValue("code")
	if !codePattern.MatchString(code) {
		writeError(w, http.StatusBadRequest, "code must look like 1234-5678")
		return
	}
	var version int64
	err := s.db.QueryRow(`SELECT version FROM sync_groups WHERE code = $1`, code).Scan(&version)
	if err == sql.ErrNoRows {
		writeError(w, http.StatusNotFound, "unknown code")
		return
	}
	if err != nil {
		writeError(w, http.StatusInternalServerError, "database error")
		return
	}
	writeJSON(w, http.StatusOK, map[string]int64{"version": version})
}

func (s *Server) putSync(w http.ResponseWriter, r *http.Request) {
	code := r.PathValue("code")
	if !codePattern.MatchString(code) {
		writeError(w, http.StatusBadRequest, "code must look like 1234-5678")
		return
	}
	body := http.MaxBytesReader(w, r.Body, 50<<20) // 50 MB ceiling (images inside)
	defer body.Close()
	var payload json.RawMessage
	if err := json.NewDecoder(body).Decode(&payload); err != nil {
		writeError(w, http.StatusBadRequest, "body must be JSON")
		return
	}
	var version int64
	err := s.db.QueryRow(`
		INSERT INTO sync_groups (code, version, data, updated_at)
		VALUES ($1, 1, $2, now())
		ON CONFLICT (code) DO UPDATE
		SET version = sync_groups.version + 1, data = $2, updated_at = now()
		RETURNING version`, code, payload).Scan(&version)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "database error")
		return
	}
	writeJSON(w, http.StatusOK, map[string]int64{"version": version})
}

func newCode() (string, error) {
	var b [8]byte
	if _, err := rand.Read(b[:]); err != nil {
		return "", err
	}
	digits := make([]byte, 8)
	for i, v := range b {
		digits[i] = '0' + v%10
	}
	return fmt.Sprintf("%s-%s", digits[:4], digits[4:]), nil
}

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

func writeError(w http.ResponseWriter, status int, message string) {
	writeJSON(w, status, map[string]string{"error": message})
}