Turbine/pkg/gateway/main.go

package main

import (
	"fmt"
	"os"
	"strings"

	"git.solsynth.dev/goatworks/turbine/pkg/shared/registrar"

	"github.com/gofiber/fiber/v3"
	"github.com/gofiber/fiber/v3/middleware/proxy"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"
	"github.com/spf13/viper"
)

var serviceDiscovery *registrar.ServiceDiscovery

func init() {
	log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr})
}

func main() {
	log.Info().Msg("Starting Turbine Gateway...")
	viper.SetConfigName("settings")
	viper.AddConfigPath(".")
	viper.AddConfigPath("/etc/turbine/")
	viper.SetConfigType("toml")

	log.Info().Msg("Reading configuration...")
	if err := viper.ReadInConfig(); err != nil {
		log.Fatal().Err(err).Msg("Failed to read config file...")
	}
	log.Info().Msg("Configuration loaded.")

	// Setup service discovery
	etcdEndpoints := viper.GetStringSlice("etcd.endpoints")
	if len(etcdEndpoints) == 0 {
		log.Fatal().Msg("etcd.endpoints not configured in settings.toml")
	}

	if viper.GetBool("etcd.insecure") {
		log.Info().Msg("Using insecure transport for etcd")
		for i, ep := range etcdEndpoints {
			if !strings.HasPrefix(ep, "http://") && !strings.HasPrefix(ep, "https://") {
				etcdEndpoints[i] = "http://" + ep
			}
		}
	}

	log.Info().Strs("endpoints", etcdEndpoints).Msg("Connecting to etcd...")

	var err error
	serviceDiscovery, err = registrar.NewServiceDiscovery(etcdEndpoints)
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to create service discovery client")
	}
	log.Info().Msg("Service discovery client created.")

	log.Info().Msg("Fetching initial service list...")
	if err := serviceDiscovery.Start(); err != nil {
		log.Fatal().Err(err).Msg("Failed to start service discovery")
	}
	log.Info().Msg("Service discovery started.")

	app := fiber.New(fiber.Config{
		ServerHeader: "Turbine Gateway",
		BodyLimit:    2147483647, // 2GB
	})

	// Health check and status
	app.Get("/", func(c fiber.Ctx) error {
		return c.JSON(fiber.Map{
			"status":   "running",
			"services": serviceDiscovery.GetServiceRoutes(),
		})
	})

	// Handle route overrides first
	routeOverrides := viper.GetStringMapString("routes")
	for from, to := range routeOverrides {
		toParts := strings.SplitN(strings.TrimPrefix(to, "/"), "/", 2)
		if len(toParts) < 1 {
			log.Warn().Str("from", from).Str("to", to).Msg("Invalid route override config")
			continue
		}
		serviceName := toParts[0]
		var remainingPath string
		if len(toParts) > 1 {
			remainingPath = toParts[1]
		}

		log.Info().Str("from", from).Str("to", to).Msg("Applying route override")

		app.Use(from, func(c fiber.Ctx) error {
			instance, err := serviceDiscovery.GetNextInstance(serviceName)
			if err != nil {
				log.Error().Err(err).Str("service", serviceName).Msg("Failed to get service instance for override")
				return c.Status(fiber.StatusServiceUnavailable).SendString(err.Error())
			}

			targetURL := fmt.Sprintf("http://%s/%s", instance, remainingPath)

			originalPath := strings.TrimPrefix(c.Path(), from)
			if originalPath != "" {
				targetURL = fmt.Sprintf("%s%s", targetURL, originalPath)
			}

			if len(c.Request().URI().QueryString()) > 0 {
				targetURL = fmt.Sprintf("%s?%s", targetURL, string(c.Request().URI().QueryString()))
			}

			log.Info().Str("from", c.Path()).Str("to", targetURL).Msg("Forwarding with override")

			if err := proxy.Do(c, targetURL); err != nil {
				return err
			}
			c.Response().SetStatusCode(fiber.StatusOK)
			return nil
		})
	}

	// Generic proxy handler
	app.Use("/*", func(c fiber.Ctx) error {
		path := c.Path()
		parts := strings.Split(strings.TrimPrefix(path, "/"), "/")
		if len(parts) < 1 || parts[0] == "" {
			// Let the health check handle this
			return c.Next()
		}

		serviceName := parts[0]
		remainingPath := strings.Join(parts[1:], "/")

		instance, err := serviceDiscovery.GetNextInstance(serviceName)
		if err != nil {
			log.Warn().Err(err).Str("service", serviceName).Msg("Failed to get service instance")
			return c.Status(fiber.StatusServiceUnavailable).SendString(err.Error())
		}

		targetURL := fmt.Sprintf("http://%s/%s", instance, remainingPath)
		if len(c.Request().URI().QueryString()) > 0 {
			targetURL = fmt.Sprintf("%s?%s", targetURL, string(c.Request().URI().QueryString()))
		}

		log.Info().Str("from", path).Str("to", targetURL).Msg("Forwarding request")

		if err := proxy.Do(c, targetURL); err != nil {
			return err
		}

		c.Response().SetStatusCode(fiber.StatusOK)
		return nil
	})

	listenAddr := viper.GetString("listen")
	if listenAddr == "" {
		listenAddr = ":8080" // default
	}
	log.Info().Msg("Gateway is listening on " + listenAddr)

	err = app.Listen(listenAddr, fiber.ListenConfig{
		DisableStartupMessage: true,
	})
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to start server...")
	}
}