main.go

// Copyright 2026 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"context"
	"flag"
	"fmt"
	"math"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"github.com/moby/buildkit/client"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"google.golang.org/grpc"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/klog"

	"github.com/tsuru/deploy-agent/pkg/build"
	"github.com/tsuru/deploy-agent/pkg/build/buildkit"
	"github.com/tsuru/deploy-agent/pkg/build/buildkit/autodiscovery"
	buildpb "github.com/tsuru/deploy-agent/pkg/build/grpc_build_v1"
	"github.com/tsuru/deploy-agent/pkg/health"
	"github.com/tsuru/deploy-agent/pkg/repository"
)

const (
	DefaultServerMaxRecvMsgSize = 4 * (1 << 30) // 4 GiB
	DefaultServerMaxSendMsgSize = math.MaxInt32 // int32 max length
)

var cfg struct {
	BuildkitAddress                                           string
	BuildkitTmpDir                                            string
	BuildKitAutoDiscoveryKubernetesPodSelector                string
	BuildKitAutoDiscoveryKubernetesNamespace                  string
	BuildKitAutoDiscoveryKubernetesLeasePrefix                string
	BuildKitAutoDiscoveryStatefulset                          string
	KubernetesConfig                                          string
	RemoteRepositoryPath                                      string
	BuildKitAutoDiscoveryTimeout                              time.Duration
	BuildKitAutoDiscoveryKubernetesPort                       int
	Port                                                      int
	MetricsPort                                               int
	ServerMaxRecvMsgSize                                      int
	ServerMaxSendMsgSize                                      int
	BuildKitAutoDiscoveryScaleGracefulPeriod                  time.Duration
	BuildKitAutoDiscovery                                     bool
	BuildKitAutoDiscoveryKubernetesSetTsuruAppLabels          bool
	BuildKitAutoDiscoveryKubernetesUseSameNamespaceAsTsuruApp bool
	DisableCache                                              bool

	// BuildKitDetectCPUArch could be use with caution only on local development
	// environments where the developer is sure about the architecture
	// of the BuildKit server and the host machine are the same.
	// This flag is NOT recommended for production environments.
	BuildKitDetectCPUArch bool
}

func main() {
	klog.InitFlags(flag.CommandLine)

	flag.IntVar(&cfg.Port, "port", 8080, "Server TCP port")
	flag.IntVar(&cfg.MetricsPort, "metrics-port", 9090, "Metrics server TCP port")
	flag.IntVar(&cfg.ServerMaxRecvMsgSize, "max-receiving-message-size", DefaultServerMaxRecvMsgSize, "Max message size in bytes that server can receive")
	flag.IntVar(&cfg.ServerMaxSendMsgSize, "max-sending-message-size", DefaultServerMaxSendMsgSize, "Max message size in bytes that server can send")

	flag.StringVar(&cfg.KubernetesConfig, "kubeconfig", getEnvOrDefault("KUBECONFIG", ""), "Path to kubeconfig file")

	flag.StringVar(&cfg.BuildkitAddress, "buildkit-addr", getEnvOrDefault("BUILDKIT_HOST", ""), "Buildkit server address")
	flag.StringVar(&cfg.BuildkitTmpDir, "buildkit-tmp-dir", os.TempDir(), "Directory path to store temp files during container image builds")

	flag.StringVar(&cfg.RemoteRepositoryPath, "remote-repository-path", getEnvOrDefault("REMOTE_REPOSITORY_PATH", ""), "Remote image repository providers config path")

	flag.BoolVar(&cfg.BuildKitAutoDiscovery, "buildkit-autodiscovery", false, "Whether should dynamically discover the BuildKit service based on Tsuru app (if any)")
	flag.DurationVar(&cfg.BuildKitAutoDiscoveryTimeout, "buildkit-autodiscovery-timeout", (5 * time.Minute), "Max duration to discover an available BuildKit")
	flag.StringVar(&cfg.BuildKitAutoDiscoveryKubernetesPodSelector, "buildkit-autodiscovery-kubernetes-pod-selector", "", "Label selector of BuildKit's pods on Kubernetes")
	flag.StringVar(&cfg.BuildKitAutoDiscoveryKubernetesNamespace, "buildkit-autodiscovery-kubernetes-namespace", "", "Namespace of BuildKit's pods on Kubernetes")
	flag.StringVar(&cfg.BuildKitAutoDiscoveryKubernetesLeasePrefix, "buildkit-autodiscovery-kubernetes-lease-prefix", "deploy-agent", "Prefix name for Lease resources")
	flag.IntVar(&cfg.BuildKitAutoDiscoveryKubernetesPort, "buildkit-autodiscovery-kubernetes-port", 80, "TCP port number which BuldKit's service is listening")
	flag.BoolVar(&cfg.BuildKitAutoDiscoveryKubernetesSetTsuruAppLabels, "buildkit-autodiscovery-kubernetes-set-tsuru-app-labels", false, "Whether should set the Tsuru app labels in the selected BuildKit pod")
	flag.BoolVar(&cfg.BuildKitAutoDiscoveryKubernetesUseSameNamespaceAsTsuruApp, "buildkit-autodiscovery-kubernetes-use-same-namespace-as-tsuru-app", false, "Whether should look for BuildKit in the Tsuru app's namespace")
	flag.StringVar(&cfg.BuildKitAutoDiscoveryStatefulset, "buildkit-autodiscovery-scale-statefulset", "", "Name of statefulset of buildkit that scale from zero")
	flag.DurationVar(&cfg.BuildKitAutoDiscoveryScaleGracefulPeriod, "buildkit-autodiscovery-scale-graceful-period", (2 * time.Hour), "how long time after a build to retain buildkit running")

	flag.BoolVar(&cfg.DisableCache, "disable-cache", false, "Disable BuildKit cache during container image builds")
	flag.BoolVar(&cfg.BuildKitDetectCPUArch, "buildkit-detect-cpu-arch", getBoolEnvOrDefault("BUILDKIT_DETECT_CPU_ARCH", false), "Whether to detect CPU architecture of the host machine and use it to pass to BuildKit")

	flag.Parse()

	l, err := net.Listen("tcp", fmt.Sprintf(":%d", cfg.Port))
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to listen: %v", err)
		os.Exit(1)
	}

	bk, err := newBuildKit()
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to create BuildKit: %v", err)
		os.Exit(1)
	}
	defer bk.Close()

	serverOpts := []grpc.ServerOption{
		grpc.MaxRecvMsgSize(cfg.ServerMaxRecvMsgSize),
		grpc.MaxSendMsgSize(cfg.ServerMaxSendMsgSize),
	}

	s := grpc.NewServer(serverOpts...)
	buildpb.RegisterBuildServer(s, build.NewServer(bk))
	healthpb.RegisterHealthServer(s, health.NewServer())

	go startMetricsServer(cfg.MetricsPort)
	go handleGracefulTermination(s)

	fmt.Println("Starting gRPC server at", l.Addr().String())

	if err := s.Serve(l); err != nil {
		fmt.Fprintln(os.Stderr, "failed to run gRPC server:", err)
		os.Exit(1)
	}

	fmt.Println("gRPC server terminated")
}

func startMetricsServer(port int) {
	http.Handle("/metrics", promhttp.Handler())
	addr := fmt.Sprintf(":%d", port)
	fmt.Println("Starting metrics server at", addr)
	server := &http.Server{
		Addr:              addr,
		ReadTimeout:       5 * time.Second,
		ReadHeaderTimeout: 5 * time.Second,
		WriteTimeout:      10 * time.Second,
	}
	if err := server.ListenAndServe(); err != nil {
		fmt.Fprintf(os.Stderr, "failed to start metrics server: %v\n", err)
		os.Exit(1)
	}
}

func handleGracefulTermination(s *grpc.Server) {
	defer func() {
		fmt.Fprintln(os.Stdout, "Received termination signal. Terminating gRPC server...")
		s.GracefulStop()
	}()

	stop := make(chan os.Signal, 1)
	signal.Notify(stop, os.Interrupt, syscall.SIGTERM)
	<-stop
}

func getEnvOrDefault(env, def string) string {
	if envvar, found := os.LookupEnv(env); found {
		return envvar
	}

	return def
}

func getBoolEnvOrDefault(env string, def bool) bool {
	if envvar, found := os.LookupEnv(env); found {
		v, _ := strconv.ParseBool(envvar)
		return v
	}

	return def
}

func newBuildKit() (*buildkit.BuildKit, error) {
	opts := buildkit.BuildKitOptions{
		TempDir:                      cfg.BuildkitTmpDir,
		DiscoverBuildKitClientForApp: cfg.BuildKitAutoDiscovery,
		DisableCache:                 cfg.DisableCache,
		DetectCPUArch:                cfg.BuildKitDetectCPUArch,
	}

	var c *client.Client

	if cfg.BuildkitAddress != "" {
		bc, err := client.New(context.Background(), cfg.BuildkitAddress, client.WithFailFast())
		if err != nil {
			return nil, fmt.Errorf("failed to create buildkit client: %w", err)
		}

		c = bc
	}

	if cfg.RemoteRepositoryPath != "" {
		repositoryData, err := os.ReadFile(cfg.RemoteRepositoryPath)
		if err != nil {
			return nil, err
		}
		opts.RemoteRepository, err = repository.NewRemoteRepository(repositoryData)
		if err != nil {
			return nil, fmt.Errorf("failed to handle remote repository cfg: %w", err)
		}
	}

	b := buildkit.NewBuildKit(c, opts)

	if cfg.BuildKitAutoDiscovery {
		config, err := clientcmd.BuildConfigFromFlags("", cfg.KubernetesConfig)
		if err != nil {
			return nil, err
		}

		cs, err := kubernetes.NewForConfig(config)
		if err != nil {
			return nil, err
		}

		dcs, err := dynamic.NewForConfig(config)
		if err != nil {
			return nil, err
		}

		kdopts := autodiscovery.KubernertesDiscoveryOptions{
			Timeout:               cfg.BuildKitAutoDiscoveryTimeout,
			PodSelector:           cfg.BuildKitAutoDiscoveryKubernetesPodSelector,
			Namespace:             cfg.BuildKitAutoDiscoveryKubernetesNamespace,
			Port:                  cfg.BuildKitAutoDiscoveryKubernetesPort,
			SetTsuruAppLabel:      cfg.BuildKitAutoDiscoveryKubernetesSetTsuruAppLabels,
			UseSameNamespaceAsApp: cfg.BuildKitAutoDiscoveryKubernetesUseSameNamespaceAsTsuruApp,
			LeasePrefix:           cfg.BuildKitAutoDiscoveryKubernetesLeasePrefix,
			Statefulset:           cfg.BuildKitAutoDiscoveryStatefulset,
			ScaleGracefulPeriod:   cfg.BuildKitAutoDiscoveryScaleGracefulPeriod,
		}

		return b.WithKubernetesDiscovery(cs, dcs, kdopts), nil
	}

	return b, nil
}