silvermobi/servers/grpcapi/server/silvermobi-grpc-server.go

package grpcserver

import (
	"context"
	"net"

	"git.coopgo.io/coopgo-apps/silvermobi/handler"
	grpcproto "git.coopgo.io/coopgo-apps/silvermobi/servers/grpcapi/proto"
	"github.com/golang-jwt/jwt/v4"
	grpcmiddleware "github.com/grpc-ecosystem/go-grpc-middleware"
	grpcauth "github.com/grpc-ecosystem/go-grpc-middleware/auth"
	grpcctxtags "github.com/grpc-ecosystem/go-grpc-middleware/tags"
	"github.com/rs/zerolog/log"
	"github.com/spf13/viper"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/status"
)

type contextKey string

var (
	contextKeyUser = contextKey("user")
)

func NoAuth(method string) bool {
	noAuthMethods := []string{
		"/SilvermobiGRPC/ForgetAccount",
		"/SilvermobiGRPC/UpdatePassword",
		"/SilvermobiGRPC/AuthRegister",
		"/SilvermobiGRPC/AuthLogin",
		"/SilvermobiGRPC/GeoAutocomplete",
		"/SilvermobiGRPC/GeoRouteWithReturn",
		"/SilvermobiGRPC/GeoRoute",
	}

	for _, m := range noAuthMethods {
		if method == m {
			return true
		}
	}
	return false
}

func UnaryAuthServerInterceptor(authFunc grpcauth.AuthFunc) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {

		var (
			newCtx context.Context
			err    error
		)

		print(info.FullMethod)

		if NoAuth(info.FullMethod) {
			return handler(ctx, req)
		}

		if newCtx, err = authFunc(ctx); err != nil {
			return nil, err
		}

		return handler(newCtx, req)
	}
}

func StreamAuthServerInterceptor(authFunc grpcauth.AuthFunc) grpc.StreamServerInterceptor {
	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo,
		handler grpc.StreamHandler) error {

		var (
			newCtx context.Context
			err    error
		)

		if NoAuth(info.FullMethod) {
			wrapped := grpcmiddleware.WrapServerStream(stream)
			wrapped.WrappedContext = stream.Context()

			return handler(srv, wrapped)
		}

		if newCtx, err = authFunc(stream.Context()); err != nil {
			return err
		}

		wrapped := grpcmiddleware.WrapServerStream(stream)
		wrapped.WrappedContext = newCtx

		return handler(srv, wrapped)
	}
}

type SilvermobiGRPCService struct {
	Config  *viper.Viper
	Handler *handler.SilverMobiHandler
	grpcproto.UnimplementedSilvermobiGRPCServer
}

type SolidarityService struct {
	Config           *viper.Viper
	Handler          *handler.SilverMobiHandler
	SolidarityClient grpcproto.SolidarityServiceClient
	grpcproto.UnimplementedSolidarityServiceServer
}

func NewSolidarityService(cfg *viper.Viper, handler *handler.SilverMobiHandler) SolidarityService {
	solidarityServiceAddress := cfg.GetString("solidarity_service.address")

	conn, err := grpc.NewClient(solidarityServiceAddress, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		log.Fatal().Err(err)
	}

	solidarityClient := grpcproto.NewSolidarityServiceClient(conn)

	return SolidarityService{
		Config:           cfg,
		Handler:          handler,
		SolidarityClient: solidarityClient,
	}
}

func NewSilvermobiGRPCService(cfg *viper.Viper, handler *handler.SilverMobiHandler) SilvermobiGRPCService {
	return SilvermobiGRPCService{
		Config:  cfg,
		Handler: handler,
	}
}

func Run(done chan error, cfg *viper.Viper, handler *handler.SilverMobiHandler) {
	var (
		address   = cfg.GetString("services.external.grpc.ip") + ":" + cfg.GetString("services.external.grpc.port")
		jwtSecret = cfg.GetString("identification.local.jwt_secret")
	)

	log.Info().Msg("GRPC server on " + address)

	server := grpc.NewServer(
		grpc.StreamInterceptor(grpcmiddleware.ChainStreamServer(
			grpcctxtags.StreamServerInterceptor(),
			StreamAuthServerInterceptor(GRPCAuthFunc(jwtSecret)),
		)),
		grpc.UnaryInterceptor(grpcmiddleware.ChainUnaryServer(
			grpcctxtags.UnaryServerInterceptor(),
			UnaryAuthServerInterceptor(GRPCAuthFunc(jwtSecret)),
		)),
	)

	solidarityService := NewSolidarityService(cfg, handler)

	silvermobiService := NewSilvermobiGRPCService(cfg, handler)
	grpcproto.RegisterSilvermobiGRPCServer(server, silvermobiService)
	grpcproto.RegisterSolidarityServiceServer(server, &solidarityService)
	l, err := net.Listen("tcp", address)
	if err != nil {
		log.Fatal().Err(err)
	}

	if err = server.Serve(l); err != nil {
		log.Error().Err(err).Msg("gRPC service ended")
		done <- err
	}
}

func GRPCAuthFunc(jwtKey string) grpcauth.AuthFunc {
	return func(ctx context.Context) (context.Context, error) {

		var token *jwt.Token

		tokenString, err := grpcauth.AuthFromMD(ctx, "bearer")

		if err != nil {
			return nil, err
		}
		claims := jwt.MapClaims{}

		token, err = jwt.ParseWithClaims(tokenString, claims, func(token *jwt.Token) (interface{}, error) {
			return []byte(jwtKey), nil
		})

		if err != nil || !token.Valid {
			return nil, status.Errorf(codes.Unauthenticated, "Invalid or expired token")
		}

		ctx = context.WithValue(ctx, contextKeyUser, claims["sub"].(string))

		return ctx, nil
	}
}