zooid/cmd/relay/main.go

package main

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"
	"zooid/zooid"
)

func main() {
	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)

	shutdown := make(chan os.Signal, 1)
	signal.Notify(shutdown, syscall.SIGINT, syscall.SIGTERM)

	port := zooid.Env("PORT")
	apiHost := zooid.Env("API_HOST")
	apiWhitelist := zooid.Env("API_WHITELIST")
	configDir := zooid.Env("CONFIG")

	// Create the main handler
	mainHandler := http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			instance, exists := zooid.Dispatch(r.Host)
			if exists {
				instance.Relay.ServeHTTP(w, r)
			} else {
				http.Error(w, "Not Found", http.StatusNotFound)
			}
		},
	)

	// Wrap with API handler if API_HOST is configured
	var handler http.Handler = mainHandler
	if apiHost != "" && apiWhitelist != "" {
		apiHandler := zooid.NewAPIHandler(apiWhitelist, configDir)
		handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// Check if this request is for the API host
			if r.Host == apiHost {
				apiHandler.ServeHTTP(w, r)
			} else {
				mainHandler.ServeHTTP(w, r)
			}
		})
	}

	srv := &http.Server{
		Addr:    fmt.Sprintf(":%s", port),
		Handler: handler,
	}

	go func() {
		log.Printf("running on :%s\n", port)
		if err := srv.ListenAndServe(); err != http.ErrServerClosed {
			log.Printf("HTTP server error: %v\n", err)
		}
	}()

	go zooid.Start()

	<-shutdown

	log.Println("\nShutting down gracefully...")

	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer shutdownCancel()

	if err := srv.Shutdown(shutdownCtx); err != nil {
		log.Printf("HTTP server shutdown error: %v\n", err)
	}
}