Merge branch 'main' into harshbhat/toggle-rootKey

unkeyed · Sep 4, 2024 · 6897023 · 6897023
2 parents 596f273 + 5748780
commit 6897023
Show file tree

Hide file tree

Showing 74 changed files with 3,880 additions and 2,271 deletions.
diff --git a/Taskfile.yml b/Taskfile.yml
@@ -12,7 +12,23 @@ tasks:
     deps: [down, build]
     cmds:
       - docker compose -f ./deployment/docker-compose.yaml up -d
+
+
   seed:
+    cmds:
+      - task: migrate-db
+      - task: migrate-clickhouse
+
+  migrate-clickhouse:
+    env:
+      GOOSE_DRIVER: clickhouse
+      GOOSE_DBSTRING: "tcp://127.0.0.1:9000"
+      GOOSE_MIGRATION_DIR: ./apps/agent/pkg/clickhouse/schema
+    cmds:
+      - goose up
+
+
+  migrate-db:
     env:
       DRIZZLE_DATABASE_URL: "mysql://unkey:password@localhost:3306/unkey"
     dir: internal/db

diff --git a/apps/agent/Dockerfile b/apps/agent/Dockerfile
@@ -16,7 +16,6 @@ COPY --from=builder /go/src/github.com/unkeyed/unkey/apps/agent/bin/unkey .
 COPY --from=builder /go/src/github.com/unkeyed/unkey/apps/agent/config.production.json .
 COPY --from=builder /go/src/github.com/unkeyed/unkey/apps/agent/config.staging.json .
 COPY --from=builder /go/src/github.com/unkeyed/unkey/apps/agent/config.docker.json .
+COPY --from=builder /go/src/github.com/unkeyed/unkey/apps/agent/pkg/openapi/openapi.json ./pkg/openapi/openapi.json
 
 CMD [ "/usr/local/bin/unkey", "agent"]
-
-
diff --git a/apps/agent/Taskfile.yml b/apps/agent/Taskfile.yml
@@ -10,8 +10,8 @@ tasks:
       - go vet ./...
   test:
     cmds:
-      - go test -cover -json -failfast ./... | tparse -all -progress 
-    
+      - go test -cover -json -failfast ./... | tparse -all -progress
+
   build:
     cmds:
       - go build -o unkey ./cmd/main.go
@@ -20,7 +20,19 @@ tasks:
     cmds:
       - go install github.com/amit-davidson/Chronos/cmd/chronos
       - ~/go/bin/chronos --file=./cmd/main.go --mod={{.TASKFILE_DIR}}
-  
+
   lint:
     cmds:
-      - golangci-lint run
+      - golangci-lint run
+
+  migrate:
+    cmds:
+      - goose -dir=./pkg/clickhouse/schema clickhouse "tcp://127.0.0.1:9000" up
+
+  generate:
+    cmds:
+      - go get github.com/oapi-codegen/oapi-codegen/v2/cmd/oapi-codegen
+      - mkdir -p ./pkg/openapi
+      - go run github.com/oapi-codegen/oapi-codegen/v2/cmd/oapi-codegen --config=./pkg/openapi/config.yaml ./pkg/openapi/openapi.json
+
+      - buf generate
diff --git a/apps/agent/cmd/agent/agent.go b/apps/agent/cmd/agent/agent.go
@@ -14,13 +14,7 @@ import (
 	"github.com/Southclaws/fault"
 	"github.com/Southclaws/fault/fmsg"
 	"github.com/unkeyed/unkey/apps/agent/pkg/api"
-	v1Liveness "github.com/unkeyed/unkey/apps/agent/pkg/api/routes/v1_liveness"
-	v1RatelimitCommitLease "github.com/unkeyed/unkey/apps/agent/pkg/api/routes/v1_ratelimit_commitLease"
-	v1RatelimitMultiRatelimit "github.com/unkeyed/unkey/apps/agent/pkg/api/routes/v1_ratelimit_multiRatelimit"
-	v1RatelimitRatelimit "github.com/unkeyed/unkey/apps/agent/pkg/api/routes/v1_ratelimit_ratelimit"
-	v1VaultDecrypt "github.com/unkeyed/unkey/apps/agent/pkg/api/routes/v1_vault_decrypt"
-	v1VaultEncrypt "github.com/unkeyed/unkey/apps/agent/pkg/api/routes/v1_vault_encrypt"
-	v1VaultEncryptBulk "github.com/unkeyed/unkey/apps/agent/pkg/api/routes/v1_vault_encrypt_bulk"
+	"github.com/unkeyed/unkey/apps/agent/pkg/clickhouse"
 	"github.com/unkeyed/unkey/apps/agent/pkg/cluster"
 	"github.com/unkeyed/unkey/apps/agent/pkg/config"
 	"github.com/unkeyed/unkey/apps/agent/pkg/connect"
@@ -133,69 +127,40 @@ func run(c *cli.Context) error {
 		setupHeartbeat(cfg, logger)
 	}
 
-	srv := api.New(api.Config{
-		NodeId:  cfg.NodeId,
-		Logger:  logger,
-		Metrics: m,
-	})
-
-	v1Liveness.Register(srv.HumaAPI(), srv.Services())
-
-	connectSrv, err := connect.New(connect.Config{Logger: logger, Image: cfg.Image, Metrics: m})
-	if err != nil {
-		return err
-	}
-
-	if cfg.Services.Vault != nil {
-		s3, err := storage.NewS3(storage.S3Config{
-			S3URL:             cfg.Services.Vault.S3Url,
-			S3Bucket:          cfg.Services.Vault.S3Bucket,
-			S3AccessKeyId:     cfg.Services.Vault.S3AccessKeyId,
-			S3AccessKeySecret: cfg.Services.Vault.S3AccessKeySecret,
-			Logger:            logger,
+	var ch clickhouse.Bufferer = clickhouse.NewNoop()
+	if cfg.Clickhouse != nil {
+		ch, err = clickhouse.New(clickhouse.Config{
+			URL:    cfg.Clickhouse.Url,
+			Logger: logger.With().Str("pkg", "clickhouse").Logger(),
 		})
 		if err != nil {
-			return fmt.Errorf("failed to create s3 storage: %w", err)
-		}
-		s3 = storageMiddleware.WithTracing("s3", s3)
-		srv.Vault, err = vault.New(vault.Config{
-			Logger:     logger,
-			Metrics:    m,
-			Storage:    s3,
-			MasterKeys: strings.Split(cfg.Services.Vault.MasterKeys, ","),
-		})
-		if err != nil {
-			return fmt.Errorf("failed to create vault: %w", err)
-		}
-
-		if err != nil {
-			return fmt.Errorf("failed to create vault service: %w", err)
+			return err
 		}
-
-		v1VaultEncrypt.Register(srv.HumaAPI(), srv.Services(), srv.BearerAuthFromSecret(cfg.Services.Vault.AuthToken))
-		v1VaultEncryptBulk.Register(srv.HumaAPI(), srv.Services(), srv.BearerAuthFromSecret(cfg.Services.Vault.AuthToken))
-		v1VaultDecrypt.Register(srv.HumaAPI(), srv.Services(), srv.BearerAuthFromSecret(cfg.Services.Vault.AuthToken))
-		logger.Info().Msg("started vault service")
 	}
 
-	if cfg.Services.EventRouter != nil {
-		er, err := eventrouter.New(eventrouter.Config{
-			Logger:        logger,
-			Metrics:       m,
-			BatchSize:     cfg.Services.EventRouter.Tinybird.BatchSize,
-			BufferSize:    cfg.Services.EventRouter.Tinybird.BufferSize,
-			FlushInterval: time.Duration(cfg.Services.EventRouter.Tinybird.FlushInterval) * time.Second,
-			Tinybird:      tinybird.New("https://api.tinybird.co", cfg.Services.EventRouter.Tinybird.Token),
-			AuthToken:     cfg.Services.EventRouter.AuthToken,
-		})
-		if err != nil {
-			return err
-		}
-		srv.WithEventRouter(er)
-		if err != nil {
-			return fmt.Errorf("failed to add event router service: %w", err)
+	s3, err := storage.NewS3(storage.S3Config{
+		S3URL:             cfg.Services.Vault.S3Url,
+		S3Bucket:          cfg.Services.Vault.S3Bucket,
+		S3AccessKeyId:     cfg.Services.Vault.S3AccessKeyId,
+		S3AccessKeySecret: cfg.Services.Vault.S3AccessKeySecret,
+		Logger:            logger,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to create s3 storage: %w", err)
+	}
+	s3 = storageMiddleware.WithTracing("s3", s3)
+	v, err := vault.New(vault.Config{
+		Logger:     logger,
+		Metrics:    m,
+		Storage:    s3,
+		MasterKeys: strings.Split(cfg.Services.Vault.MasterKeys, ","),
+	})
+	if err != nil {
+		return fmt.Errorf("failed to create vault: %w", err)
+	}
 
-		}
+	if err != nil {
+		return fmt.Errorf("failed to create vault service: %w", err)
 	}
 
 	var clus cluster.Cluster
@@ -255,39 +220,65 @@ func run(c *cli.Context) error {
 			}
 		}()
 
-		err = connectSrv.AddService(connect.NewClusterServer(clus, logger))
-		if err != nil {
-			return fmt.Errorf("failed to add cluster service: %w", err)
+	}
 
-		}
+	rl, err := ratelimit.New(ratelimit.Config{
+		Logger:  logger,
+		Metrics: m,
+		Cluster: clus,
+	})
+	if err != nil {
+		logger.Fatal().Err(err).Msg("failed to create service")
 	}
 
-	if cfg.Services.Ratelimit != nil {
-		rl, err := ratelimit.New(ratelimit.Config{
-			Logger:  logger,
-			Metrics: m,
-			Cluster: clus,
+	srv, err := api.New(api.Config{
+		NodeId:     cfg.NodeId,
+		Logger:     logger,
+		Ratelimit:  rl,
+		Metrics:    m,
+		Clickhouse: ch,
+		AuthToken:  cfg.Cluster.AuthToken,
+		Vault:      v,
+	})
+	if err != nil {
+		return err
+	}
+
+	if cfg.Services.EventRouter != nil {
+		er, err := eventrouter.New(eventrouter.Config{
+			Logger:        logger,
+			Metrics:       m,
+			BatchSize:     cfg.Services.EventRouter.Tinybird.BatchSize,
+			BufferSize:    cfg.Services.EventRouter.Tinybird.BufferSize,
+			FlushInterval: time.Duration(cfg.Services.EventRouter.Tinybird.FlushInterval) * time.Second,
+			Tinybird:      tinybird.New("https://api.tinybird.co", cfg.Services.EventRouter.Tinybird.Token),
+			AuthToken:     cfg.AuthToken,
 		})
 		if err != nil {
-			logger.Fatal().Err(err).Msg("failed to create service")
+			return err
 		}
-
-		srv.Ratelimit = ratelimit.WithTracing(rl)
-
-		v1RatelimitRatelimit.Register(srv.HumaAPI(), srv.Services(), srv.BearerAuthFromSecret(cfg.Services.Ratelimit.AuthToken))
-		v1RatelimitMultiRatelimit.Register(srv.HumaAPI(), srv.Services(), srv.BearerAuthFromSecret(cfg.Services.Ratelimit.AuthToken))
-		v1RatelimitCommitLease.Register(srv.HumaAPI(), srv.Services(), srv.BearerAuthFromSecret(cfg.Services.Ratelimit.AuthToken))
-
-		err = connectSrv.AddService(connect.NewRatelimitServer(rl, logger, cfg.Services.Ratelimit.AuthToken))
+		srv.WithEventRouter(er)
 		if err != nil {
-			return fmt.Errorf("failed to add ratelimit service: %w", err)
+			return fmt.Errorf("failed to add event router service: %w", err)
+
 		}
-		logger.Info().Msg("started ratelimit service")
 	}
 
-	if cfg.Pprof != nil {
-		connectSrv.EnablePprof(cfg.Pprof.Username, cfg.Pprof.Password)
+	connectSrv, err := connect.New(connect.Config{Logger: logger, Image: cfg.Image, Metrics: m})
+	if err != nil {
+		return err
+	}
+
+	err = connectSrv.AddService(connect.NewClusterServer(clus, logger))
+	if err != nil {
+		return fmt.Errorf("failed to add cluster service: %w", err)
+
+	}
+	err = connectSrv.AddService(connect.NewRatelimitServer(rl, logger, cfg.AuthToken))
+	if err != nil {
+		return fmt.Errorf("failed to add ratelimit service: %w", err)
 	}
+	logger.Info().Msg("started ratelimit service")
 
 	go func() {
 		err := connectSrv.Listen(fmt.Sprintf(":%s", cfg.RpcPort))

diff --git a/apps/agent/cmd/main.go b/apps/agent/cmd/main.go
@@ -26,7 +26,14 @@ func main() {
 		fmt.Println()
 		fmt.Println()
 
-		fmt.Println(chain)
+		for _, e := range chain {
+			fmt.Printf(" - ")
+			if e.Location != "" {
+				fmt.Printf("%s\n", e.Location)
+				fmt.Printf("    > ")
+			}
+			fmt.Printf("%s\n", e.Message)
+		}
 		fmt.Println()
 		os.Exit(1)
 	}

diff --git a/apps/agent/config.docker.json b/apps/agent/config.docker.json
@@ -6,6 +6,7 @@
     "username": "admin",
     "password": "admin"
   },
+  "authToken": "${AUTH_TOKEN}",
   "nodeId": "${NODE_ID}",
   "logging": {},
   "cluster": {
@@ -19,11 +20,7 @@
     }
   },
   "services": {
-    "ratelimit": {
-      "authToken": "${AUTH_TOKEN}"
-    },
     "eventRouter": {
-      "authToken": "${AUTH_TOKEN}",
       "tinybird": {
         "token": "${TINYBIRD_TOKEN}",
         "batchSize": 1000,
@@ -36,12 +33,14 @@
       "s3Bucket": "${VAULT_S3_BUCKET}",
       "s3AccessKeyId": "${VAULT_S3_ACCESS_KEY_ID}",
       "s3AccessKeySecret": "${VAULT_S3_ACCESS_KEY_SECRET}",
-      "masterKeys": "${VAULT_MASTER_KEYS}",
-      "authToken": "${AUTH_TOKEN}"
+      "masterKeys": "${VAULT_MASTER_KEYS}"
     }
   },
   "prometheus": {
     "path": "/metrics",
     "port": 2112
+  },
+  "clickhouse": {
+    "url": "${CLICKHOUSE_URL}"
   }
 }
diff --git a/apps/agent/config.production.json b/apps/agent/config.production.json
@@ -10,6 +10,7 @@
   "port": "${PORT}",
   "rpcPort": "${RPC_PORT}",
   "region": "fly::${FLY_REGION}",
+  "authToken": "${AUTH_TOKEN}",
   "logging": {
     "axiom": {
       "dataset": "agent",
@@ -29,11 +30,7 @@
     }
   },
   "services": {
-    "ratelimit": {
-      "authToken": "${AUTH_TOKEN}"
-    },
     "eventRouter": {
-      "authToken": "${AUTH_TOKEN}",
       "tinybird": {
         "token": "${TINYBIRD_TOKEN}",
         "batchSize": 1000,
@@ -46,8 +43,7 @@
       "s3Bucket": "${VAULT_S3_BUCKET}",
       "s3AccessKeyId": "${VAULT_S3_ACCESS_KEY_ID}",
       "s3AccessKeySecret": "${VAULT_S3_ACCESS_KEY_SECRET}",
-      "masterKeys": "${VAULT_MASTER_KEYS}",
-      "authToken": "${AUTH_TOKEN}"
+      "masterKeys": "${VAULT_MASTER_KEYS}"
     }
   },
   "cluster": {

diff --git a/apps/agent/config.staging.json b/apps/agent/config.staging.json
@@ -10,12 +10,9 @@
   "port": "${PORT}",
   "rpcPort": "${RPC_PORT}",
   "region": "fly::${FLY_REGION}",
+  "authToken": "${AUTH_TOKEN}",
   "services": {
-    "ratelimit": {
-      "authToken": "${AUTH_TOKEN}"
-    },
     "eventRouter": {
-      "authToken": "${AUTH_TOKEN}",
       "tinybird": {
         "token": "${TINYBIRD_TOKEN}",
         "batchSize": 1000,
@@ -28,8 +25,7 @@
       "s3Bucket": "${VAULT_S3_BUCKET}",
       "s3AccessKeyId": "${VAULT_S3_ACCESS_KEY_ID}",
       "s3AccessKeySecret": "${VAULT_S3_ACCESS_KEY_SECRET}",
-      "masterKeys": "${VAULT_MASTER_KEYS}",
-      "authToken": "${AUTH_TOKEN}"
+      "masterKeys": "${VAULT_MASTER_KEYS}"
     }
   },
   "cluster": {