gitlabhq/gitlab-shell

// Package main implements the GitLab SSH daemon.
package main

import (
    "context"
    "flag"
    "os"
    "os/signal"
    "syscall"
    "time"

    "gitlab.com/gitlab-org/gitlab-shell/v14/internal/command"
    "gitlab.com/gitlab-org/gitlab-shell/v14/internal/config"
    "gitlab.com/gitlab-org/gitlab-shell/v14/internal/logger"
    "gitlab.com/gitlab-org/gitlab-shell/v14/internal/sshd"

    "gitlab.com/gitlab-org/labkit/log"
    "gitlab.com/gitlab-org/labkit/monitoring"
)

var (
    configDir = flag.String("config-dir", "", "The directory the config is in")

    // Version is the current version of gitlab-shell
    Version = "(unknown version)" // Set at build time in the Makefile
    // BuildTime signifies the time the binary was build
    BuildTime = "19700101.000000" // Set at build time in the Makefile
)

func overrideConfigFromEnvironment(cfg *config.Config) {
    if gitlabURL := os.Getenv("GITLAB_URL"); gitlabURL != "" {
        cfg.GitlabUrl = gitlabURL
    }
    if gitlabTracing := os.Getenv("GITLAB_TRACING"); gitlabTracing != "" {
        cfg.GitlabTracing = gitlabTracing
    }
    if gitlabShellSecret := os.Getenv("GITLAB_SHELL_SECRET"); gitlabShellSecret != "" {
        cfg.Secret = gitlabShellSecret
    }
    if gitlabLogFormat := os.Getenv("GITLAB_LOG_FORMAT"); gitlabLogFormat != "" {
        cfg.LogFormat = gitlabLogFormat
    }
}

func main() {
    command.CheckForVersionFlag(os.Args, Version, BuildTime)

    flag.Parse()

    cfg := new(config.Config)
    if *configDir != "" {
        var err error
        cfg, err = config.NewFromDir(*configDir)
        if err != nil {
            log.WithError(err).Fatal("failed to load configuration from specified directory")
        }
    }

    overrideConfigFromEnvironment(cfg)
    if err := cfg.IsSane(); err != nil {
        if *configDir == "" {
            log.WithError(err).Fatal("no config-dir provided, using only environment variables")
        } else {
            log.WithError(err).Fatal("configuration error")
        }
    }

    cfg.ApplyGlobalState()

    logCloser := logger.ConfigureStandalone(cfg)
    defer func() {
        if err := logCloser.Close(); err != nil {
            log.WithError(err).Fatal("Error closing logCloser")
        }
    }()
    ctx, finished := command.Setup("gitlab-sshd", cfg)
    defer finished()

    cfg.GitalyClient.InitSidechannelRegistry(ctx)

    server, err := sshd.NewServer(cfg)
    if err != nil {
        log.WithError(err).Fatal("Failed to start GitLab built-in sshd")
    }

    // Startup monitoring endpoint.
    if cfg.Server.WebListen != "" {
        startupMonitoringEndpoint(cfg, server)
    }

    ctx, cancel := context.WithCancel(ctx)
    defer cancel()

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

    gracefulShutdown(ctx, done, cfg, server, cancel)

    if err := server.ListenAndServe(ctx); err != nil {
        log.WithError(err).Fatal("GitLab built-in sshd failed to listen for new connections")
    }
}

func gracefulShutdown(ctx context.Context, done chan os.Signal, cfg *config.Config, server *sshd.Server, cancel context.CancelFunc) {
    go func() {
        sig := <-done
        signal.Reset(syscall.SIGINT, syscall.SIGTERM)

        gracePeriod := time.Duration(cfg.Server.GracePeriod)
        log.WithContextFields(ctx, log.Fields{"shutdown_timeout_s": gracePeriod.Seconds(), "signal": sig.String()}).Info("Shutdown initiated")

        if err := server.Shutdown(); err != nil {
            log.WithError(err).Fatal("Error shutting down the server")
        }

        <-time.After(gracePeriod)

        cancel()
    }()
}

func startupMonitoringEndpoint(cfg *config.Config, server *sshd.Server) {
    go func() {
        err := monitoring.Start(
            monitoring.WithListenerAddress(cfg.Server.WebListen),
            monitoring.WithBuildInformation(Version, BuildTime),
            monitoring.WithServeMux(server.MonitoringServeMux()),
        )

        log.WithError(err).Fatal("monitoring service raised an error")
    }()
}