[management] Prevent JWT reuse during peer login (#6002)

management/internals/shared/grpc/server.go

@@ -14,6 +14,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	jwtv5 "github.com/golang-jwt/jwt/v5"
 	pb "github.com/golang/protobuf/proto" // nolint
 	"github.com/golang/protobuf/ptypes/timestamp"
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/realip"
@@ -67,6 +68,7 @@ type Server struct {
 	appMetrics     telemetry.AppMetrics
 	peerLocks      sync.Map
 	authManager    auth.Manager
+	sessionStore   *auth.SessionStore
 
 	logBlockedPeers          bool
 	blockPeersWithSameConfig bool
@@ -98,6 +100,7 @@ func NewServer(
 	integratedPeerValidator integrated_validator.IntegratedValidator,
 	networkMapController network_map.Controller,
 	oAuthConfigProvider idp.OAuthConfigProvider,
+	sessionStore *auth.SessionStore,
 ) (*Server, error) {
 	if appMetrics != nil {
 		// update gauge based on number of connected peers which is equal to open gRPC streams
@@ -140,6 +143,7 @@ func NewServer(
 		integratedPeerValidator:  integratedPeerValidator,
 		networkMapController:     networkMapController,
 		oAuthConfigProvider:      oAuthConfigProvider,
+		sessionStore:             sessionStore,
 
 		loginFilter: newLoginFilter(),
 
@@ -535,7 +539,7 @@ func (s *Server) cancelPeerRoutinesWithoutLock(ctx context.Context, accountID st
 	log.WithContext(ctx).Debugf("peer %s has been disconnected", peer.Key)
 }
 
-func (s *Server) validateToken(ctx context.Context, jwtToken string) (string, error) {
+func (s *Server) validateToken(ctx context.Context, peerKey, jwtToken string) (string, error) {
 	if s.authManager == nil {
 		return "", status.Errorf(codes.Internal, "missing auth manager")
 	}
@@ -545,6 +549,10 @@ func (s *Server) validateToken(ctx context.Context, jwtToken string) (string, er
 		return "", status.Errorf(codes.InvalidArgument, "invalid jwt token, err: %v", err)
 	}
 
+	if err := s.claimLoginToken(ctx, peerKey, jwtToken, token); err != nil {
+		return "", err
+	}
+
 	// we need to call this method because if user is new, we will automatically add it to existing or create a new account
 	accountId, _, err := s.accountManager.GetAccountIDFromUserAuth(ctx, userAuth)
 	if err != nil {
@@ -828,6 +836,31 @@ func (s *Server) prepareLoginResponse(ctx context.Context, peer *nbpeer.Peer, ne
 	return loginResp, nil
 }
 
+func (s *Server) claimLoginToken(ctx context.Context, peerKey, jwtToken string, token *jwtv5.Token) error {
+	if s.sessionStore == nil || token == nil {
+		return nil
+	}
+
+	exp, err := token.Claims.GetExpirationTime()
+	if err != nil || exp == nil {
+		log.WithContext(ctx).Warnf("JWT has no usable exp claim for peer %s", peerKey)
+		return status.Error(codes.Unauthenticated, "jwt token has no expiration")
+	}
+
+	err = s.sessionStore.RegisterToken(ctx, jwtToken, exp.Time)
+	if err == nil {
+		return nil
+	}
+
+	if errors.Is(err, auth.ErrTokenAlreadyUsed) || errors.Is(err, auth.ErrTokenExpired) {
+		log.WithContext(ctx).Warnf("%v for peer %s", err, peerKey)
+		return status.Error(codes.Unauthenticated, err.Error())
+	}
+
+	log.WithContext(ctx).Warnf("failed to claim JWT for peer %s: %v", peerKey, err)
+	return status.Error(codes.Unavailable, "failed to claim jwt token")
+}
+
 // processJwtToken validates the existence of a JWT token in the login request, and returns the corresponding user ID if
 // the token is valid.
 //
@@ -838,7 +871,7 @@ func (s *Server) processJwtToken(ctx context.Context, loginReq *proto.LoginReque
 	if loginReq.GetJwtToken() != "" {
 		var err error
 		for i := 0; i < 3; i++ {
-			userID, err = s.validateToken(ctx, loginReq.GetJwtToken())
+			userID, err = s.validateToken(ctx, peerKey.String(), loginReq.GetJwtToken())
 			if err == nil {
 				break
 			}