github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/cmd/authn/main.go (about)

     1  // Package authn is authentication server for AIStore.
     2  /*
     3   * Copyright (c) 2018-2023, NVIDIA CORPORATION. All rights reserved.
     4   */
     5  package main
     6  
     7  import (
     8  	"flag"
     9  	"fmt"
    10  	"os"
    11  	"os/signal"
    12  	"path/filepath"
    13  	"strings"
    14  	"syscall"
    15  	"time"
    16  
    17  	"github.com/NVIDIA/aistore/api/env"
    18  	"github.com/NVIDIA/aistore/cmn"
    19  	"github.com/NVIDIA/aistore/cmn/cos"
    20  	"github.com/NVIDIA/aistore/cmn/fname"
    21  	"github.com/NVIDIA/aistore/cmn/jsp"
    22  	"github.com/NVIDIA/aistore/cmn/kvdb"
    23  	"github.com/NVIDIA/aistore/cmn/nlog"
    24  )
    25  
    26  const secretKeyPodEnv = "SECRETKEY" // via https://kubernetes.io/docs/concepts/configuration/secret
    27  
    28  var (
    29  	build     string
    30  	buildtime string
    31  
    32  	configPath string
    33  )
    34  
    35  func init() {
    36  	flag.StringVar(&configPath, "config", "", svcName+" configuration")
    37  }
    38  
    39  func logFlush() {
    40  	for {
    41  		time.Sleep(time.Minute) // TODO: must be configurable
    42  		nlog.Flush(nlog.ActNone)
    43  	}
    44  }
    45  
    46  func main() {
    47  	var configDir string
    48  	if len(os.Args) == 2 && os.Args[1] == "version" {
    49  		printVer()
    50  		os.Exit(0)
    51  	}
    52  	if len(os.Args) == 1 || (len(os.Args) == 2 && strings.Contains(os.Args[1], "help")) {
    53  		printVer()
    54  		flag.PrintDefaults()
    55  		os.Exit(0)
    56  	}
    57  	installSignalHandler()
    58  	flag.Parse()
    59  
    60  	confDirFlag := flag.Lookup("config")
    61  	if confDirFlag != nil {
    62  		configDir = confDirFlag.Value.String()
    63  	}
    64  	if configDir == "" {
    65  		configDir = os.Getenv(env.AuthN.ConfDir)
    66  	}
    67  	if configDir == "" {
    68  		cos.ExitLogf("Missing %s configuration file (to specify, use '-%s' option or '%s' environment)",
    69  			svcName, confDirFlag.Name, env.AuthN.ConfDir)
    70  	}
    71  	configPath = filepath.Join(configDir, fname.AuthNConfig)
    72  	if _, err := jsp.LoadMeta(configPath, Conf); err != nil {
    73  		cos.ExitLogf("Failed to load configuration from %q: %v", configPath, err)
    74  	}
    75  	if val := os.Getenv(secretKeyPodEnv); val != "" {
    76  		Conf.Server.Secret = val
    77  	}
    78  	if err := updateLogOptions(); err != nil {
    79  		cos.ExitLogf("Failed to set up logger: %v", err)
    80  	}
    81  	if Conf.Verbose() {
    82  		nlog.Infof("Loaded configuration from %s", configPath)
    83  	}
    84  
    85  	dbPath := filepath.Join(configDir, fname.AuthNDB)
    86  	driver, err := kvdb.NewBuntDB(dbPath)
    87  	if err != nil {
    88  		cos.ExitLogf("Failed to init local database: %v", err)
    89  	}
    90  	mgr, err := newMgr(driver)
    91  	if err != nil {
    92  		cos.ExitLogf("Failed to init manager: %v", err)
    93  	}
    94  
    95  	nlog.Infof("Version %s (build %s)\n", cmn.VersionAuthN+"."+build, buildtime)
    96  
    97  	go logFlush()
    98  
    99  	srv := newServer(mgr)
   100  	err = srv.Run()
   101  
   102  	nlog.Flush(nlog.ActExit)
   103  	cos.Close(mgr.db)
   104  	if err != nil {
   105  		cos.ExitLogf("Server failed: %v", err)
   106  	}
   107  }
   108  
   109  func updateLogOptions() error {
   110  	if err := cos.CreateDir(Conf.Log.Dir); err != nil {
   111  		return fmt.Errorf("failed to create log dir %q, err: %v", Conf.Log.Dir, err)
   112  	}
   113  	nlog.SetLogDirRole(Conf.Log.Dir, "auth")
   114  	return nil
   115  }
   116  
   117  func installSignalHandler() {
   118  	c := make(chan os.Signal, 1)
   119  	signal.Notify(c, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
   120  	go func() {
   121  		<-c
   122  		os.Exit(0)
   123  	}()
   124  }
   125  
   126  func printVer() {
   127  	fmt.Printf("version %s (build %s)\n", cmn.VersionAuthN+"."+build, buildtime)
   128  }