backend/cloudprint-backend/internal/grpc/interceptor/auth.go

package interceptor

import (
	"context"
	"strings"

	"github.com/cloudprint/cloudprint-backend/pkg/jwt"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

const (
	AuthorizationHeader = "authorization"
	BearerPrefix        = "Bearer "
	UserIDKey           = "user_id"
	UserTypeKey         = "user_type"
)

type AuthInterceptor struct {
	jwtManager *jwt.JWTManager
}

func NewAuthInterceptor(jwtManager *jwt.JWTManager) *AuthInterceptor {
	return &AuthInterceptor{jwtManager: jwtManager}
}

// UnaryServerInterceptor returns a unary server interceptor that handles authentication
func (i *AuthInterceptor) UnaryServerInterceptor() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		// Skip auth for certain methods
		if isPublicMethod(info.FullMethod) {
			return handler(ctx, req)
		}

		claims, err := i.extractAndValidateToken(ctx)
		if err != nil {
			return nil, status.Errorf(codes.Unauthenticated, "invalid token: %v", err)
		}

		// Add user info to context
		ctx = context.WithValue(ctx, UserIDKey, claims.UserID)
		ctx = context.WithValue(ctx, UserTypeKey, claims.UserType)

		return handler(ctx, req)
	}
}

// StreamServerInterceptor returns a streaming server interceptor that handles authentication
func (i *AuthInterceptor) StreamServerInterceptor() grpc.StreamServerInterceptor {
	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		// Skip auth for certain methods
		if isPublicMethod(info.FullMethod) {
			return handler(srv, ss)
		}

		claims, err := i.extractAndValidateToken(ss.Context())
		if err != nil {
			return status.Errorf(codes.Unauthenticated, "invalid token: %v", err)
		}

		// Add user info to context
		ctx := context.WithValue(ss.Context(), UserIDKey, claims.UserID)
		ctx = context.WithValue(ctx, UserTypeKey, claims.UserType)

		wrapped := &serverStreamWrapper{
			ServerStream: ss,
			ctx:          ctx,
		}

		return handler(srv, wrapped)
	}
}

func (i *AuthInterceptor) extractAndValidateToken(ctx context.Context) (*jwt.Claims, error) {
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		return nil, status.Error(codes.Unauthenticated, "missing metadata")
	}

	authHeader := md.Get(AuthorizationHeader)
	if len(authHeader) == 0 {
		return nil, status.Error(codes.Unauthenticated, "missing authorization header")
	}

	tokenString := strings.TrimPrefix(authHeader[0], BearerPrefix)
	if tokenString == authHeader[0] {
		return nil, status.Error(codes.Unauthenticated, "invalid authorization format")
	}

	return i.jwtManager.ValidateToken(tokenString)
}

func isPublicMethod(method string) bool {
	publicMethods := map[string]bool{
		"/cloudprint.AuthService/Login":            true,
		"/cloudprint.AuthService/RefreshToken":     true,
		"/cloudprint.AdminService/Login":           true,
		"/cloudprint.PayService/HandlePayCallback": true,
	}
	return publicMethods[method]
}

// serverStreamWrapper wraps grpc.ServerStream to override Context method
type serverStreamWrapper struct {
	grpc.ServerStream
	ctx context.Context
}

func (w *serverStreamWrapper) Context() context.Context {
	return w.ctx
}

// GetUserIDFromContext extracts user ID from context
func GetUserIDFromContext(ctx context.Context) (int32, bool) {
	userID, ok := ctx.Value(UserIDKey).(int32)
	return userID, ok
}

// GetUserTypeFromContext extracts user type from context
func GetUserTypeFromContext(ctx context.Context) (string, bool) {
	userType, ok := ctx.Value(UserTypeKey).(string)
	return userType, ok
}