Token authorization is ready for MVP

Signed-off-by: Nikolai Rodionov <iam@allanger.xyz>

internal/controllers/authorization.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"slices"
 	"strings"
 	"time"
 
@@ -22,6 +23,8 @@ type TokenType string
 const (
 	TokenTypeAccess  TokenType = "access"
 	TokenTypeRefresh TokenType = "refresh"
+	TokenAudToken    string    = "token"
+	TokenAudWeb      string    = "web"
 )
 
 var (
@@ -33,6 +36,7 @@ type Claims struct {
 	UserID    string    `json:"user_id"`
 	TokenID   string    `json:"token_id"`
 	TokenType TokenType `json:"token_type"`
+	Scope     string    `json:"scope,omitempty"`
 	jwt.RegisteredClaims
 }
 
@@ -56,6 +60,13 @@ func NewAuthController(jwtSecret []byte, accessTTL, refreshTTL time.Duration, re
 	}
 }
 
+type JWTData struct {
+	UserID    string
+	TokenType TokenType
+	TokenAud  string
+	Scope     string
+}
+
 // Write claims into context
 func (a *AuthController) WithClaims(ctx context.Context, claims *Claims) context.Context {
 	return context.WithValue(ctx, claimsContextKey, claims)
@@ -87,15 +98,43 @@ func (a *AuthController) AuthInterceptorFN(ctx context.Context) (context.Context
 		}
 	}
 
+	// If it's a cli token, we need to check the scope
+	if slices.Contains(claims.Audience, TokenAudToken) {
+		currentMethod, ok := grpc.Method(ctx)
+		if !ok {
+			return nil, errors.New("unknown method")
+		}
+
+		scopeMap := map[string][]string{}
+		if err := json.Unmarshal([]byte(claims.Scope), &scopeMap); err != nil {
+			return nil, ErrServerError
+		}
+		allowed := isAllowed(scopeMap, currentMethod)
+		if !allowed {
+			return nil, errors.New("not authorized")
+		}
+	}
+
 	ctx = a.WithClaims(ctx, claims)
 	return ctx, nil
 }
 
+func isAllowed(scope map[string][]string, currentMethod string) bool {
+	for service, methods := range scope {
+		for _, method := range methods {
+			if fmt.Sprintf("/%s/%s", service, method) == currentMethod {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 // Generate JWT token
-func (a *AuthController) GenerateToken(userID string, tokenType TokenType) (token, tokenID string, err error) {
+func (a *AuthController) GenerateToken(data *JWTData) (token, tokenID string, err error) {
 	var expiresAt time.Time
 	notBefore := time.Now()
-	switch tokenType {
+	switch data.TokenType {
 	case TokenTypeAccess:
 		expiresAt = time.Now().Add(a.accessTTL)
 	case TokenTypeRefresh:
@@ -103,25 +142,25 @@ func (a *AuthController) GenerateToken(userID string, tokenType TokenType) (toke
 	default:
 		return "", "", ErrUnknownTokenType
 	}
-	if tokenType != TokenTypeAccess && tokenType != TokenTypeRefresh {
-		return "", "", ErrUnknownTokenType
-	}
 
 	tokenID = uuid.New().String()
+
 	claims := Claims{
-		UserID:    userID,
+		UserID:    data.UserID,
 		TokenID:   tokenID,
-		TokenType: tokenType,
+		TokenType: data.TokenType,
+		Scope:     data.Scope,
 		RegisteredClaims: jwt.RegisteredClaims{
 			Issuer:    "",
-			Subject:   "",
-			Audience:  jwt.ClaimStrings{},
+			Subject:   data.UserID,
+			Audience:  jwt.ClaimStrings{data.TokenAud},
 			ExpiresAt: jwt.NewNumericDate(expiresAt),
 			NotBefore: jwt.NewNumericDate(notBefore),
 			IssuedAt:  jwt.NewNumericDate(time.Now()),
-			ID:        userID,
+			ID:        tokenID,
 		},
 	}
+
 	tokenJwt := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
 	token, err = tokenJwt.SignedString(a.jwtSecret)
 	if err != nil {

internal/controllers/authorization_test.go

@@ -16,16 +16,26 @@ var (
 )
 
 func TestGenerateInvalidTokenType(t *testing.T) {
+	data := &controllers.JWTData{
+		UserID:    testUserID,
+		TokenType: "invalid_type",
+	}
+
 	authCtrl := controllers.NewAuthController([]byte("test"), testAccessTTL, testRefreshTTL, nil)
-	token, _, err := authCtrl.GenerateToken(testUserID, "invalid_type")
+
+	token, _, err := authCtrl.GenerateToken(data)
 	assert.Equal(t, "", token)
 	assert.ErrorIs(t, controllers.ErrUnknownTokenType, err)
 }
 
 func TestGenerateValidateAccessToken(t *testing.T) {
+	data := &controllers.JWTData{
+		UserID:    testUserID,
+		TokenType: controllers.TokenTypeAccess,
+	}
 	authCtrl := controllers.NewAuthController([]byte("test"), testAccessTTL, testRefreshTTL, nil)
 	now := time.Now()
-	token, _, err := authCtrl.GenerateToken(testUserID, controllers.TokenTypeAccess)
+	token, _, err := authCtrl.GenerateToken(data)
 	assert.NoError(t, err)
 	assert.NotEmpty(t, token)
 
@@ -40,9 +50,13 @@ func TestGenerateValidateAccessToken(t *testing.T) {
 }
 
 func TestGenerateValidateRefreshToken(t *testing.T) {
+	data := &controllers.JWTData{
+		UserID:    testUserID,
+		TokenType: controllers.TokenTypeRefresh,
+	}
 	authCtrl := controllers.NewAuthController([]byte("test"), testAccessTTL, testRefreshTTL, nil)
 	now := time.Now()
-	token, _, err := authCtrl.GenerateToken(testUserID, controllers.TokenTypeRefresh)
+	token, _, err := authCtrl.GenerateToken(data)
 	assert.NoError(t, err)
 	assert.NotEmpty(t, token)
 

internal/controllers/tokens.go

@@ -330,22 +330,27 @@ func shouldSkip(s string, rules []rule) bool {
 	return false
 }
 
-func (ctrl *TokenController) AuthenticateWithToken(ctx context.Context, token string) (map[string][]string, error) {
+type TokenAuthResult struct {
+	UserID string
+	Scope  string
+}
+
+func (ctrl *TokenController) AuthenticateWithToken(ctx context.Context, token string) (*TokenAuthResult, error) {
 	log := logger.FromContext(ctx)
 	log.V(2).Info("Authenticating with a token")
 
 	query := `
-	SELECT scopes, expires_at, revoked_at
+	SELECT user_id, scopes, expires_at, revoked_at
 	FROM tokens
 	WHERE token_hash = $1`
 
+	var userID string
 	var expiresAt sql.NullTime
 	var revokedAt sql.NullTime
-	var scopes string
-	fmt.Println(hashSHA256(token))
-	fmt.Println(hashSHA256(token))
+	var scope string
 	if err := ctrl.DB.QueryRowContext(ctx, query, hashSHA256(token)).Scan(
-		&scopes,
+		&userID,
+		&scope,
 		&expiresAt,
 		&revokedAt,
 	); err != nil {
@@ -356,20 +361,20 @@ func (ctrl *TokenController) AuthenticateWithToken(ctx context.Context, token st
 		return nil, ErrServerError
 	}
 
-	if !revokedAt.Valid {
+	if revokedAt.Valid {
 		return nil, ErrBadToken
 	}
 
-	if expiresAt.Time.After(time.Now()) {
+	if expiresAt.Time.Before(time.Now()) {
 		return nil, ErrBadToken
 	}
 
-	scopesMap := map[string][]string{}
-
-	if err := json.Unmarshal([]byte(scopes), scopesMap); err != nil {
-		return nil, ErrServerError
+	result := &TokenAuthResult{
+		UserID: userID,
+		Scope:  scope,
 	}
-	return scopesMap, nil
+
+	return result, nil
 }
 
 func hashSHA256(s string) string {