JWT Authentication (#106)

* auth config + genKey

* jwt auth & tests

* update swagger docs & re-add APIError

cmd/zep/cmd.go

@@ -14,6 +14,7 @@ var (
 
 	cfgFile     string
 	showVersion bool
+	generateKey bool
 )
 
 var cmd = &cobra.Command{
@@ -24,8 +25,9 @@ var cmd = &cobra.Command{
 
 func init() {
 	cmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default config.yaml)")
-	// add cobra switch to print version number
 	cmd.PersistentFlags().BoolVarP(&showVersion, "version", "v", false, "print version number")
+	cmd.PersistentFlags().
+		BoolVarP(&generateKey, "generate-token", "g", false, "generate a new JWT token")
 }
 
 // Execute executes the root cobra command.

cmd/zep/run.go

@@ -8,6 +8,8 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/getzep/zep/pkg/auth"
+
 	"github.com/oiime/logrusbun"
 	"github.com/sirupsen/logrus"
 	"github.com/uptrace/bun"
@@ -28,17 +30,15 @@ const (
 
 // run is the entrypoint for the zep server
 func run() {
-	if showVersion {
-		fmt.Println(VersionString)
-		os.Exit(0)
-	}
-	log.Infof("Starting zep server version %s", VersionString)
-
 	cfg, err := config.LoadConfig(cfgFile)
 	if err != nil {
 		log.Fatalf("Error configuring Zep: %s", err)
 	}
 
+	handleCLIOptions(cfg)
+
+	log.Infof("Starting zep server version %s", VersionString)
+
 	config.SetLogLevel(cfg)
 	appState := NewAppState(cfg)
 
@@ -69,6 +69,18 @@ func NewAppState(cfg *config.Config) *models.AppState {
 	return appState
 }
 
+// handleCLIOptions handles CLI options that don't require the server to run
+func handleCLIOptions(cfg *config.Config) {
+	if showVersion {
+		fmt.Println(VersionString)
+		os.Exit(0)
+	}
+	if generateKey {
+		fmt.Println(auth.GenerateJWT(cfg))
+		os.Exit(0)
+	}
+}
+
 // initializeMemoryStore initializes the memory store based on the config file / ENV
 func initializeMemoryStore(appState *models.AppState) {
 	if appState.Config.MemoryStore.Type == "" {

config.yaml

@@ -26,10 +26,15 @@ memory_store:
     dsn: "postgres://postgres:postgres@localhost:5432/?sslmode=disable"
 server:
   port: 8000
+auth:
+  # Set to true to enable authentication
+  required: false
+  # Do not use this secret in production. The ZEP_AUTH_SECRET environment variable should be
+  # set to a cryptographically secure secret. See the Zep docs for details.
+  secret: "do-not-use-this-secret-in-production"
 data:
   #  PurgeEvery is the period between hard deletes, in minutes.
   #  If set to 0 or undefined, hard deletes will not be performed.
   purge_every: 60
 log:
-  level: "info"
-
+  level: "info"

config/config.go

@@ -13,6 +13,12 @@ import (
 // We're bootstrapping so avoid any imports from other packages
 var log = logrus.New()
 
+// EnvVars is a set of secrets that should be stored in the environment, not config file
+var EnvVars = map[string]string{
+	"llm.openai_api_key": "ZEP_OPENAI_API_KEY",
+	"auth.secret":        "ZEP_AUTH_SECRET",
+}
+
 // LoadConfig loads the config file and ENV variables into a Config struct
 func LoadConfig(configFile string) (*Config, error) {
 	if configFile != "" {
@@ -39,9 +45,9 @@ func LoadConfig(configFile string) (*Config, error) {
 	// Environment variables take precedence over config file
 	loadDotEnv()
 
-	err := viper.BindEnv("llm.openai_api_key", "ZEP_OPENAI_API_KEY")
-	if err != nil {
-		log.Fatalf("Error binding environment variable: %s", err)
+	// Bind environment variables to config keys
+	for key, envVar := range EnvVars {
+		bindEnv(key, envVar)
 	}
 
 	var cfg Config
@@ -63,6 +69,14 @@ func loadDotEnv() {
 	}
 }
 
+// bindEnv binds an environment variable to a config key
+func bindEnv(key string, envVar string) {
+	err := viper.BindEnv(key, envVar)
+	if err != nil {
+		log.Fatalf("Error binding environment variable: %s", err)
+	}
+}
+
 // SetLogLevel sets the log level based on the config file. Defaults to INFO if not set or invalid
 func SetLogLevel(cfg *Config) {
 	level, err := logrus.ParseLevel(cfg.Log.Level)

config/models.go

@@ -10,6 +10,7 @@ type Config struct {
 	MemoryStore MemoryStoreConfig `mapstructure:"memory_store"`
 	Server      ServerConfig      `mapstructure:"server"`
 	Log         LogConfig         `mapstructure:"log"`
+	Auth        AuthConfig        `mapstructure:"auth"`
 	DataConfig  DataConfig        `mapstructure:"data"`
 }
 
@@ -46,6 +47,11 @@ type LogConfig struct {
 	Level string `mapstructure:"level"`
 }
 
+type AuthConfig struct {
+	Secret   string `mapstructure:"secret"`
+	Required bool   `mapstructure:"required"`
+}
+
 type DataConfig struct {
 	// PurgeEvery is the period between hard deletes, in minutes.
 	// If set to 0, hard deletes will not be performed.