feat: ip whitelisting (#130)

* feat: ip whitelisting * refactor(key_verify.go): remove unnecessary logging statement refactor(server.go): reformat fiber.Config initialization for better readability
unkeyed · Jul 8, 2023 · a9cd5e7 · a9cd5e7 · vercel · Jul 8, 2023
1 parent e76f272
commit a9cd5e7
Show file tree

Hide file tree

Showing 22 changed files with 501 additions and 65 deletions.
diff --git a/apps/api/Taskfile.yml b/apps/api/Taskfile.yml
@@ -49,7 +49,7 @@ tasks:
       - docker push ghcr.io/chronark/{{.IMAGE}}:{{.VERSION}}
 
   debug:
-    - go test -v -json -shuffle=on --race ./... -run TestRefresh | gotestfmt
+    - go test -v -json -shuffle=on --race ./... -run TestVerifyKey_WithIpWhitelist | gotestfmt
   introspect:
     # https://github.com/xo/xo
     - xo schema $XO_DSN -o ./pkg/database/models --go-pkg=models
diff --git a/apps/api/cmd/api/main.go b/apps/api/cmd/api/main.go
@@ -122,26 +122,35 @@ func main() {
 	}
 	db = databaseMiddleware.WithTracing(db, tracer)
 
-	c := cache.New[entities.Key](cache.Config[entities.Key]{
+	keyCache := cache.New[entities.Key](cache.Config[entities.Key]{
 		Fresh:             time.Minute,
 		Stale:             time.Minute * 15,
 		RefreshFromOrigin: db.GetKeyByHash,
 		Logger:            logger,
 	})
-	c = cacheMiddleware.WithTracing[entities.Key](c, tracer)
-	c = cacheMiddleware.WithLogging[entities.Key](c, logger)
+	keyCache = cacheMiddleware.WithTracing[entities.Key](keyCache, tracer)
+	keyCache = cacheMiddleware.WithLogging[entities.Key](keyCache, logger)
+
+	apiCache := cache.New[entities.Api](cache.Config[entities.Api]{
+		Fresh:             time.Minute,
+		Stale:             time.Minute * 15,
+		RefreshFromOrigin: db.GetApi,
+		Logger:            logger,
+	})
+	apiCache = cacheMiddleware.WithTracing[entities.Api](apiCache, tracer)
+	apiCache = cacheMiddleware.WithLogging[entities.Api](apiCache, logger)
 
 	k.RegisterOnKeyEvent(func(ctx context.Context, e kafka.KeyEvent) error {
 		logger.Info("evicting key from cache", zap.String("keyId", e.Key.Id), zap.String("keyHash", e.Key.Hash))
-		c.Remove(context.Background(), e.Key.Hash)
+		keyCache.Remove(context.Background(), e.Key.Hash)
 
 		if e.Type == kafka.KeyCreated || e.Type == kafka.KeyUpdated {
 			logger.Info("fetching key from origin", zap.String("keyId", e.Key.Id), zap.String("keyHash", e.Key.Hash))
 			key, err := db.GetKeyById(ctx, e.Key.Id)
 			if err != nil {
 				return fmt.Errorf("unable to get key by id: %s: %w", e.Key.Id, err)
 			}
-			c.Set(ctx, key.Hash, key)
+			keyCache.Set(ctx, key.Hash, key)
 		}
 
 		return nil
@@ -151,7 +160,8 @@ func main() {
 
 	srv := server.New(server.Config{
 		Logger:            logger,
-		Cache:             c,
+		KeyCache:          keyCache,
+		ApiCache:          apiCache,
 		Database:          db,
 		Ratelimit:         fastRatelimit,
 		GlobalRatelimit:   consistentRatelimit,

diff --git a/apps/api/fly.toml → apps/api/fly.prod.toml b/apps/api/fly.toml → apps/api/fly.prod.toml
diff --git a/apps/api/pkg/database/api_get.go b/apps/api/pkg/database/api_get.go
@@ -21,9 +21,5 @@ func (db *database) GetApi(ctx context.Context, apiId string) (entities.Api, err
 	if api == nil {
 		return entities.Api{}, ErrNotFound
 	}
-	return entities.Api{
-		Id:          api.ID,
-		Name:        api.Name,
-		WorkspaceId: api.WorkspaceID,
-	}, nil
+	return apiModelToEntity(api), nil
 }
diff --git a/apps/api/pkg/database/api_test.go b/apps/api/pkg/database/api_test.go
@@ -0,0 +1,69 @@
+package database
+
+import (
+	"context"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/unkeyed/unkey/apps/api/pkg/entities"
+	"github.com/unkeyed/unkey/apps/api/pkg/logging"
+	"github.com/unkeyed/unkey/apps/api/pkg/uid"
+)
+
+func TestApiSetGet_Simple(t *testing.T) {
+	ctx := context.Background()
+	db, err := New(Config{
+		Logger:    logging.NewNoopLogger(),
+		PrimaryUs: os.Getenv("DATABASE_DSN"),
+	})
+
+	require.NoError(t, err)
+
+	api := entities.Api{
+		Id:          uid.Api(),
+		Name:        "test",
+		WorkspaceId: uid.Workspace(),
+	}
+
+	err = db.CreateApi(ctx, api)
+	require.NoError(t, err)
+
+	found, err := db.GetApi(ctx, api.Id)
+	require.NoError(t, err)
+
+	require.Equal(t, api.Id, found.Id)
+	require.Equal(t, api.Name, found.Name)
+	require.Equal(t, api.WorkspaceId, found.WorkspaceId)
+	require.Equal(t, 0, len(found.IpWhitelist))
+
+}
+
+func TestApiSetGet_WithIpWhitelist(t *testing.T) {
+	ctx := context.Background()
+	db, err := New(Config{
+		Logger:    logging.NewNoopLogger(),
+		PrimaryUs: os.Getenv("DATABASE_DSN"),
+	})
+
+	require.NoError(t, err)
+
+	api := entities.Api{
+		Id:          uid.Api(),
+		Name:        "test",
+		WorkspaceId: uid.Workspace(),
+		IpWhitelist: []string{"1.1.1.1", "2.2.2.2"},
+	}
+
+	err = db.CreateApi(ctx, api)
+	require.NoError(t, err)
+
+	found, err := db.GetApi(ctx, api.Id)
+	require.NoError(t, err)
+
+	require.Equal(t, api.Id, found.Id)
+	require.Equal(t, api.Name, found.Name)
+	require.Equal(t, api.WorkspaceId, found.WorkspaceId)
+	require.Equal(t, api.IpWhitelist, found.IpWhitelist)
+
+}
diff --git a/apps/api/pkg/database/conversion.go b/apps/api/pkg/database/conversion.go
@@ -4,6 +4,7 @@ import (
 	"database/sql"
 	"encoding/json"
 	"fmt"
+	"strings"
 
 	"github.com/unkeyed/unkey/apps/api/pkg/database/models"
 	"github.com/unkeyed/unkey/apps/api/pkg/entities"
@@ -123,15 +124,20 @@ func apiEntityToModel(a entities.Api) *models.API {
 		ID:          a.Id,
 		Name:        a.Name,
 		WorkspaceID: a.WorkspaceId,
+		IPWhitelist: sql.NullString{String: strings.Join(a.IpWhitelist, ","), Valid: len(a.IpWhitelist) > 0},
 	}
 
 }
 
 func apiModelToEntity(model *models.API) entities.Api {
-	return entities.Api{
+	a := entities.Api{
 		Id:          model.ID,
 		Name:        model.Name,
 		WorkspaceId: model.WorkspaceID,
 	}
+	if model.IPWhitelist.Valid {
+		a.IpWhitelist = strings.Split(model.IPWhitelist.String, ",")
+	}
+	return a
 
 }
diff --git a/apps/api/pkg/database/conversion_test.go b/apps/api/pkg/database/conversion_test.go
@@ -0,0 +1,82 @@
+package database
+
+import (
+	"database/sql"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/unkeyed/unkey/apps/api/pkg/database/models"
+	"github.com/unkeyed/unkey/apps/api/pkg/entities"
+	"github.com/unkeyed/unkey/apps/api/pkg/uid"
+)
+
+func Test_apiEntityToModel(t *testing.T) {
+	e := entities.Api{
+		Id:          uid.Api(),
+		Name:        "test",
+		WorkspaceId: uid.Workspace(),
+	}
+
+	m := apiEntityToModel(e)
+
+	require.Equal(t, e.Id, m.ID)
+	require.Equal(t, e.Name, m.Name)
+	require.Equal(t, e.WorkspaceId, m.WorkspaceID)
+	require.Equal(t, false, m.IPWhitelist.Valid)
+	require.Equal(t, "", m.IPWhitelist.String)
+
+}
+
+func Test_apiModelToEntity(t *testing.T) {
+
+	m := &models.API{
+		ID:          uid.Api(),
+		Name:        "test",
+		WorkspaceID: uid.Workspace(),
+	}
+
+	e := apiModelToEntity(m)
+
+	require.Equal(t, m.ID, e.Id)
+	require.Equal(t, m.Name, e.Name)
+	require.Equal(t, m.WorkspaceID, e.WorkspaceId)
+	require.Equal(t, 0, len(e.IpWhitelist))
+}
+
+func Test_apiEntityToModel_WithIpWithlist(t *testing.T) {
+	e := entities.Api{
+		Id:          uid.Api(),
+		Name:        "test",
+		WorkspaceId: uid.Workspace(),
+		IpWhitelist: []string{"1.1.1.1", "2.2.2.2"},
+	}
+
+	m := apiEntityToModel(e)
+
+	require.Equal(t, e.Id, m.ID)
+	require.Equal(t, e.Name, m.Name)
+	require.Equal(t, e.WorkspaceId, m.WorkspaceID)
+	require.Equal(t, true, m.IPWhitelist.Valid)
+	require.Equal(t, "1.1.1.1,2.2.2.2", m.IPWhitelist.String)
+
+}
+
+func Test_apiModelToEntity_WithIpWithlist(t *testing.T) {
+
+	m := &models.API{
+		ID:          uid.Api(),
+		Name:        "test",
+		WorkspaceID: uid.Workspace(),
+		IPWhitelist: sql.NullString{
+			Valid:  true,
+			String: "1.1.1.1,2.2.2.2",
+		},
+	}
+
+	e := apiModelToEntity(m)
+
+	require.Equal(t, m.ID, e.Id)
+	require.Equal(t, m.Name, e.Name)
+	require.Equal(t, m.WorkspaceID, e.WorkspaceId)
+	require.Equal(t, []string{"1.1.1.1", "2.2.2.2"}, e.IpWhitelist)
+}
diff --git a/apps/api/pkg/database/models/api.xo.go b/apps/api/pkg/database/models/api.xo.go
diff --git a/apps/api/pkg/database/workspace_get.go b/apps/api/pkg/database/workspace_get.go
@@ -16,8 +16,5 @@ func (db *database) GetWorkspace(ctx context.Context, workspaceId string) (entit
 	if workspace == nil {
 		return entities.Workspace{}, fmt.Errorf("unable to find workspace %s in db", workspaceId)
 	}
-	return entities.Workspace{
-		Id:   workspace.ID,
-		Name: workspace.Name,
-	}, nil
+	return workspaceModelToEntity(workspace), nil
 }
diff --git a/apps/api/pkg/entities/entities.go b/apps/api/pkg/entities/entities.go
@@ -28,6 +28,7 @@ type Api struct {
 	Id          string
 	Name        string
 	WorkspaceId string
+	IpWhitelist []string
 }
 
 type Workspace struct {

diff --git a/apps/api/pkg/server/api_get.go b/apps/api/pkg/server/api_get.go
@@ -14,9 +14,10 @@ type GetApiRequest struct {
 }
 
 type GetApiResponse struct {
-	Id          string `json:"id"`
-	Name        string `json:"name"`
-	WorkspaceId string `json:"workspaceId"`
+	Id          string   `json:"id"`
+	Name        string   `json:"name"`
+	WorkspaceId string   `json:"workspaceId"`
+	IpWhitelist []string `json:"ipWhitelist,omitempty"`
 }
 
 func (s *Server) getApi(c *fiber.Ctx) error {
@@ -79,5 +80,6 @@ func (s *Server) getApi(c *fiber.Ctx) error {
 		Id:          api.Id,
 		Name:        api.Name,
 		WorkspaceId: api.WorkspaceId,
+		IpWhitelist: api.IpWhitelist,
 	})
 }