proxy-pool/internal/controller/runtime/runtime.go

package runtime

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/http"
	"strings"

	"github.com/proxy-pool/proxy-pool/internal/config"
	"github.com/proxy-pool/proxy-pool/internal/controller/admin"
	"github.com/proxy-pool/proxy-pool/internal/controller/distribution"
	"github.com/proxy-pool/proxy-pool/internal/platform/httpsecurity"
	"github.com/proxy-pool/proxy-pool/internal/platform/httpserver"
)

const (
	defaultDistributionBodyLimitBytes int64 = 64 << 10
	defaultAdminBodyLimitBytes        int64 = 8 << 10
)

var ErrInvalidRuntime = errors.New("invalid controller HTTP runtime")

type Dependencies struct {
	Extractor    distribution.Extractor
	Readiness    distribution.ReadinessChecker
	AdminService admin.Service
}

type Options struct {
	DistributionBodyLimitBytes int64
	AdminBodyLimitBytes        int64
	HTTP                       httpserver.Options
}

type Listeners struct {
	Distribution net.Listener
	Admin        net.Listener
}

type Runtime struct {
	distributionEnabled bool
	distributionAddress string
	distributionHandler http.Handler
	adminEnabled        bool
	adminAddress        string
	adminHandler        http.Handler
	httpOptions         httpserver.Options
}

func New(cfg *config.Config, dependencies Dependencies, options Options) (*Runtime, error) {
	if cfg == nil || (!cfg.Distribution.Enabled && !cfg.Admin.Enabled) {
		return nil, ErrInvalidRuntime
	}
	distributionBodyLimit, adminBodyLimit, err := resolveBodyLimits(options)
	if err != nil {
		return nil, err
	}
	result := &Runtime{httpOptions: options.HTTP}

	if cfg.Distribution.Enabled {
		if strings.TrimSpace(cfg.Distribution.Listen) == "" || dependencies.Extractor == nil || dependencies.Readiness == nil {
			return nil, ErrInvalidRuntime
		}
		protection, err := httpsecurity.BuildFromListener(
			cfg.Distribution.Listener,
			cfg.Distribution.ClientIdentification.Mode,
			httpsecurity.APIAuthSemantics,
		)
		if err != nil {
			return nil, fmt.Errorf("%w: build distribution protection: %v", ErrInvalidRuntime, err)
		}
		handler, err := distribution.NewHandler(distribution.Config{BodyLimitBytes: distributionBodyLimit}, distribution.Dependencies{
			Extractor: dependencies.Extractor,
			Identity:  protection,
			Readiness: dependencies.Readiness,
		})
		if err != nil {
			return nil, fmt.Errorf("%w: %v", ErrInvalidRuntime, err)
		}
		result.distributionEnabled = true
		result.distributionAddress = cfg.Distribution.Listen
		result.distributionHandler = handler
	}

	if cfg.Admin.Enabled {
		if strings.TrimSpace(cfg.Admin.Listen) == "" || dependencies.AdminService == nil {
			return nil, ErrInvalidRuntime
		}
		protection, err := httpsecurity.BuildFromListener(
			cfg.Admin,
			httpsecurity.ClientSourceIP,
			httpsecurity.APIAuthSemantics,
		)
		if err != nil {
			return nil, fmt.Errorf("%w: build admin protection: %v", ErrInvalidRuntime, err)
		}
		handler, err := admin.NewHandler(dependencies.AdminService, protection, admin.Options{MaxBodyBytes: adminBodyLimit})
		if err != nil {
			return nil, fmt.Errorf("%w: %v", ErrInvalidRuntime, err)
		}
		result.adminEnabled = true
		result.adminAddress = cfg.Admin.Listen
		result.adminHandler = handler
	}
	return result, nil
}

func (runtime *Runtime) Run(ctx context.Context) error {
	if runtime == nil || ctx == nil {
		return ErrInvalidRuntime
	}
	bindings := make([]httpserver.Binding, 0, 2)
	if runtime.distributionEnabled {
		bindings = append(bindings, httpserver.Binding{
			Name: "distribution", Address: runtime.distributionAddress, Handler: runtime.distributionHandler,
		})
	}
	if runtime.adminEnabled {
		bindings = append(bindings, httpserver.Binding{
			Name: "admin", Address: runtime.adminAddress, Handler: runtime.adminHandler,
		})
	}
	if err := httpserver.ListenAndServe(ctx, runtime.httpOptions, bindings...); err != nil {
		return fmt.Errorf("run controller HTTP runtime: %w", err)
	}
	return nil
}

func (runtime *Runtime) Serve(ctx context.Context, listeners Listeners) error {
	if runtime == nil || ctx == nil ||
		runtime.distributionEnabled != (listeners.Distribution != nil) ||
		runtime.adminEnabled != (listeners.Admin != nil) {
		return ErrInvalidRuntime
	}
	endpoints := make([]httpserver.Endpoint, 0, 2)
	if runtime.distributionEnabled {
		endpoints = append(endpoints, httpserver.Endpoint{
			Name: "distribution", Listener: listeners.Distribution, Handler: runtime.distributionHandler,
		})
	}
	if runtime.adminEnabled {
		endpoints = append(endpoints, httpserver.Endpoint{
			Name: "admin", Listener: listeners.Admin, Handler: runtime.adminHandler,
		})
	}
	if err := httpserver.Serve(ctx, runtime.httpOptions, endpoints...); err != nil {
		return fmt.Errorf("serve controller HTTP runtime: %w", err)
	}
	return nil
}

func resolveBodyLimits(options Options) (int64, int64, error) {
	if options.DistributionBodyLimitBytes < 0 || options.AdminBodyLimitBytes < 0 {
		return 0, 0, ErrInvalidRuntime
	}
	distributionLimit := options.DistributionBodyLimitBytes
	if distributionLimit == 0 {
		distributionLimit = defaultDistributionBodyLimitBytes
	}
	adminLimit := options.AdminBodyLimitBytes
	if adminLimit == 0 {
		adminLimit = defaultAdminBodyLimitBytes
	}
	return distributionLimit, adminLimit, nil
}