gRPC Authentication Interceptor

package middleware

import (
    "context"
    "errors"
    "sort"
    "strings"

    "github.com/rs/zerolog/log"
    "google.golang.org/grpc"
    "google.golang.org/grpc/codes"
    "google.golang.org/grpc/metadata"
    "google.golang.org/grpc/status"

    "github.com/BrosSquad/ts-1-chat-app/backend/services/auth"
)

type AuthConfig struct {
    Allowed sort.StringSlice
    Blocked sort.StringSlice
}

func authCheck(ctx context.Context, token, method string, config *AuthConfig, tokenService auth.TokenService) error {
    blockedLen := len(config.Blocked)
    allowedLen := len(config.Allowed)

    // Both lists are empty, everything should go through
    if blockedLen == 0 && allowedLen == 0 {
        return nil
    }

    if blockedLen != 0 {
        idx := config.Blocked.Search(method)

        // Found in blacklist
        if idx != blockedLen {
            return tokenService.Verify(ctx, token)
        }

        // If not found in blacklist, then it is allowed to pass
        return nil
    }

    if allowedLen != 0 {
        idx := config.Allowed.Search(method)

        // Found in whitelist
        if idx != allowedLen {
            return nil
        }
    }

    // If not found in whitelist, then token must be checked
    return tokenService.Verify(ctx, token)
}

func extractToken(ctx context.Context) (string, error) {
    md, ok := metadata.FromIncomingContext(ctx)

    if !ok {
        return "", errors.New("metadata not found")
    }

    header := md.Get("authorization")

    if len(header) != 1 {
        return "", errors.New("authorization header not found")
    }

    typeAndToken := strings.SplitN(header[0], " ", 1)

    if len(typeAndToken) != 2 {
        return "", errors.New("authorization token invalid format: $TOKEN_TYPE$ $TOKEN$")
    }

    return typeAndToken[1], nil
}

func UnaryAuth(config AuthConfig, tokenService auth.TokenService) grpc.UnaryServerInterceptor {
    if len(config.Blocked) != 0 && len(config.Allowed) != 0 {
        log.Fatal().Msg("Use either whitelist or blacklist")
    }

    return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
        method, exist := grpc.Method(ctx)

        if !exist {
            return nil, status.Error(codes.Unauthenticated, "unauthenticated")
        }

        log.Debug().Str("method", method).Msg("Unary Request method name")

        token, err := extractToken(ctx)

        if err != nil {
            return nil, status.Error(codes.Unauthenticated, err.Error())
        }

        if err := authCheck(ctx, token, method, &config, tokenService); err != nil {
            return nil, status.Error(codes.Unauthenticated, "unauthenticated")
        }

        return handler(ctx, req)
    }
}

func StreamAuth(config AuthConfig, tokenService auth.TokenService) grpc.StreamServerInterceptor {
    return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
        method, exist := grpc.MethodFromServerStream(ss)

        if !exist {
            return status.Error(codes.Unauthenticated, "unauthenticated")
        }

        log.Debug().Str("method", method).Msg("Stream Request method name")

        token, err := extractToken(ss.Context())

        if err != nil {
            return status.Error(codes.Unauthenticated, err.Error())
        }

        if err := authCheck(ss.Context(), token, method, &config, tokenService); err != nil {
            return status.Error(codes.Unauthenticated, "unauthenticated")
        }

        return handler(srv, ss)
    }
}