Rewrite access tokens to sqlc

Author: jakubno

packages/db/queries/builds/validate_build.sql

@@ -0,0 +1,6 @@
+-- name: ValidateEnvBuilds :one
+SELECT at.user_id FROM envs e
+JOIN users_teams ut on ut.team_id = e.team_id
+JOIN access_tokens at on at.user_id = ut.user_id
+WHERE at.access_token_hash = $1
+  AND e.id = @template_id;

packages/db/queries/delete_access_tokens.sql.go

@@ -0,0 +1,3 @@
+-- name: Test_DeleteAccessToken :exec
+DELETE FROM "public"."access_tokens"
+WHERE user_id = $1;

packages/db/queries/tests/delete_access_tokens.sql

@@ -1,7 +1,7 @@
 version: "2"
 sql:
   - engine: "postgresql"
-    queries: "queries"
+    queries: "queries/**"
     schema: "migrations"
     gen:
       go:

packages/db/queries/validate_build.sql.go

@@ -4,25 +4,24 @@ go 1.24.7
 
 replace github.com/e2b-dev/infra/packages/shared v0.0.0 => ../shared
 
+replace github.com/e2b-dev/infra/packages/db v0.0.0 => ../db
+
 require (
+	github.com/e2b-dev/infra/packages/db v0.0.0
 	github.com/e2b-dev/infra/packages/shared v0.0.0
 	github.com/jellydator/ttlcache/v3 v3.4.0
 )
 
 require (
-	ariga.io/atlas v0.15.0 // indirect
-	entgo.io/ent v0.12.5 // indirect
-	github.com/agext/levenshtein v1.2.3 // indirect
-	github.com/apparentlymart/go-textseg/v15 v15.0.0 // indirect
-	github.com/dchest/uniuri v1.2.0 // indirect
-	github.com/go-openapi/inflect v0.21.0 // indirect
-	github.com/go-test/deep v1.0.8 // indirect
-	github.com/google/go-cmp v0.7.0 // indirect
 	github.com/google/uuid v1.6.0 // indirect
-	github.com/hashicorp/hcl/v2 v2.19.1 // indirect
+	github.com/jackc/pgpassfile v1.0.0 // indirect
+	github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
+	github.com/jackc/pgx/v5 v5.7.4 // indirect
+	github.com/jackc/puddle/v2 v2.2.2 // indirect
 	github.com/lib/pq v1.10.9 // indirect
-	github.com/mitchellh/go-wordwrap v1.0.1 // indirect
-	github.com/zclconf/go-cty v1.14.1 // indirect
+	go.uber.org/multierr v1.11.0 // indirect
+	go.uber.org/zap v1.27.0 // indirect
+	golang.org/x/crypto v0.42.0 // indirect
 	golang.org/x/mod v0.27.0 // indirect
 	golang.org/x/sync v0.17.0 // indirect
 	golang.org/x/text v0.29.0 // indirect

packages/db/sqlc.yaml

@@ -2,60 +2,52 @@ package auth
 
 import (
 	"context"
+	"database/sql"
 	"encoding/base64"
+	"errors"
 	"fmt"
 	"log"
 	"strings"
 
+	"github.com/e2b-dev/infra/packages/db/client"
+	"github.com/e2b-dev/infra/packages/db/queries"
 	"github.com/e2b-dev/infra/packages/shared/pkg/keys"
-	"github.com/e2b-dev/infra/packages/shared/pkg/models"
-	"github.com/e2b-dev/infra/packages/shared/pkg/models/accesstoken"
-	"github.com/e2b-dev/infra/packages/shared/pkg/models/env"
-	"github.com/e2b-dev/infra/packages/shared/pkg/models/envbuild"
-	"github.com/e2b-dev/infra/packages/shared/pkg/models/user"
 )
 
-func Validate(ctx context.Context, db *models.Client, token, envID string) (bool, error) {
+func Validate(ctx context.Context, sqlcDB *client.Client, token, envID string) (bool, error) {
 	hashedToken, err := keys.VerifyKey(keys.AccessTokenPrefix, token)
 	if err != nil {
 		return false, err
 	}
 
-	u, err := db.User.Query().Where(user.HasAccessTokensWith(accesstoken.AccessTokenHash(hashedToken))).WithTeams().Only(ctx)
+	_, err = sqlcDB.ValidateEnvBuilds(ctx, queries.ValidateEnvBuildsParams{
+		TemplateID:      envID,
+		AccessTokenHash: hashedToken,
+	})
 	if err != nil {
-		return false, err
-	}
+		if errors.Is(err, sql.ErrNoRows) {
+			return false, nil
+		}
 
-	e, err := db.Env.Query().Where(
-		env.ID(envID),
-		env.HasBuildsWith(envbuild.StatusEQ(envbuild.StatusWaiting)),
-	).Only(ctx)
-	if err != nil {
 		return false, err
 	}
 
-	for _, team := range u.Edges.Teams {
-		if team.ID == e.TeamID {
-			return true, nil
-		}
-	}
-
-	return false, nil
+	return true, nil
 }
 
-func ValidateAccessToken(ctx context.Context, db *models.Client, accessToken string) bool {
+func ValidateAccessToken(ctx context.Context, db *client.Client, accessToken string) bool {
 	hashedToken, err := keys.VerifyKey(keys.AccessTokenPrefix, accessToken)
 	if err != nil {
 		return false
 	}
 
-	exists, err := db.AccessToken.Query().Where(accesstoken.AccessTokenHash(hashedToken)).Exist(ctx)
+	_, err = db.GetUserIDFromAccessToken(ctx, hashedToken)
 	if err != nil {
 		log.Printf("Error while checking access token: %s\n", err.Error())
 		return false
 	}
 
-	return exists
+	return true
 }
 
 func ExtractAccessToken(authHeader, authType string) (string, error) {

packages/docker-reverse-proxy/go.mod

@@ -1,27 +1,28 @@
 package handlers
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"log"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
 
+	"github.com/e2b-dev/infra/packages/db/client"
 	"github.com/e2b-dev/infra/packages/docker-reverse-proxy/internal/cache"
 	"github.com/e2b-dev/infra/packages/shared/pkg/consts"
-	"github.com/e2b-dev/infra/packages/shared/pkg/db"
 )
 
 type APIStore struct {
-	db        *db.DB
+	db        *client.Client
 	AuthCache *cache.AuthCache
 	proxy     *httputil.ReverseProxy
 }
 
-func NewStore() *APIStore {
+func NewStore(ctx context.Context) *APIStore {
 	authCache := cache.New()
-	database, err := db.NewClient(3, 2)
+	database, err := client.NewClient(ctx, client.WithMaxConnections(3))
 	if err != nil {
 		log.Fatal(err)
 	}

packages/docker-reverse-proxy/go.sum

@@ -37,7 +37,7 @@ func (a *APIStore) GetToken(w http.ResponseWriter, r *http.Request) error {
 		return fmt.Errorf("error while extracting access token: %w", err)
 	}
 
-	if !auth.ValidateAccessToken(ctx, a.db.Client, accessToken) {
+	if !auth.ValidateAccessToken(ctx, a.db, accessToken) {
 		log.Printf("Invalid access token: '%s'\n", accessToken)
 
 		w.WriteHeader(http.StatusForbidden)
@@ -78,7 +78,7 @@ func (a *APIStore) GetToken(w http.ResponseWriter, r *http.Request) error {
 	}
 
 	// Validate if the user has access to the template
-	hasAccess, err := auth.Validate(ctx, a.db.Client, accessToken, templateID)
+	hasAccess, err := auth.Validate(ctx, a.db, accessToken, templateID)
 	if err != nil {
 		w.WriteHeader(http.StatusInternalServerError)