Support persistent session in IntrospectToken method of GRPC client, …

…partly PLTFRM-72444
acronis · Jan 29, 2025 · ca77c0a · ca77c0a
1 parent 094a015
commit ca77c0a
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 4 deletions.
diff --git a/idptoken/grpc_client.go b/idptoken/grpc_client.go
@@ -8,8 +8,10 @@ package idptoken
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"strconv"
+	"sync/atomic"
 	"time"
 
 	"github.com/acronis/go-appkit/log"
@@ -33,6 +35,7 @@ const DefaultGRPCClientRequestTimeout = time.Second * 30
 const (
 	grpcMetaAuthorization = "authorization"
 	grpcMetaRequestID     = "x-request-id"
+	grpcMetaSessionID     = "x-session-id"
 )
 
 // GRPCClientOpts contains options for the GRPCClient.
@@ -62,6 +65,8 @@ type GRPCClient struct {
 	reqTimeout        time.Duration
 	promMetrics       *metrics.PrometheusMetrics
 	requestIDProvider func(ctx context.Context) string
+
+	sessionID atomic.Value
 }
 
 // NewGRPCClient creates a new GRPCClient instance that communicates with the IDP token service.
@@ -126,7 +131,11 @@ func (c *GRPCClient) IntrospectToken(
 		req.ScopeFilter[i] = &pb.IntrospectionScopeFilter{ResourceNamespace: scopeFilter[i].ResourceNamespace}
 	}
 
-	ctx = metadata.AppendToOutgoingContext(ctx, grpcMetaAuthorization, makeBearerToken(accessToken))
+	if sessID := c.GetSessionID(); sessID != "" {
+		ctx = metadata.AppendToOutgoingContext(ctx, grpcMetaSessionID, sessID)
+	} else {
+		ctx = metadata.AppendToOutgoingContext(ctx, grpcMetaAuthorization, makeBearerToken(accessToken))
+	}
 	if c.requestIDProvider != nil {
 		ctx = metadata.AppendToOutgoingContext(ctx, grpcMetaRequestID, c.requestIDProvider(ctx))
 	}
@@ -135,11 +144,20 @@ func (c *GRPCClient) IntrospectToken(
 	if err := c.do(ctx, "IDPTokenService/IntrospectToken", func(ctx context.Context) error {
 		var innerErr error
 		resp, innerErr = c.client.IntrospectToken(ctx, &req)
+		if errors.Is(innerErr, ErrUnauthenticated) {
+			c.SetSessionID("")
+		}
 		return innerErr
 	}); err != nil {
 		return nil, err
 	}
 
+	if md, ok := metadata.FromIncomingContext(ctx); ok {
+		if sessionIDList := md.Get(grpcMetaSessionID); len(sessionIDList) > 0 {
+			c.SetSessionID(sessionIDList[0])
+		}
+	}
+
 	claims := jwt.DefaultClaims{
 		RegisteredClaims: jwtgo.RegisteredClaims{
 			Issuer:   resp.GetIss(),
@@ -177,6 +195,18 @@ func (c *GRPCClient) IntrospectToken(
 	}, nil
 }
 
+func (c *GRPCClient) GetSessionID() string {
+	id, ok := c.sessionID.Load().(string)
+	if !ok {
+		return ""
+	}
+	return id
+}
+
+func (c *GRPCClient) SetSessionID(id string) {
+	c.sessionID.Store(id)
+}
+
 type exchangeTokenOptions struct {
 	notRequiredIntrospection bool
 }

diff --git a/idptoken/introspector_test.go b/idptoken/introspector_test.go
@@ -162,6 +162,7 @@ func TestIntrospector_IntrospectToken(t *gotesting.T) {
 		introspectorOpts        idptoken.IntrospectorOpts
 		tokenToIntrospect       string
 		accessToken             string
+		sessionID               string
 		expectedResult          idptoken.IntrospectionResult
 		checkError              func(t *gotesting.T, err error)
 		expectedHTTPSrvCalled   bool

diff --git a/internal/testing/server_token_introspector_mock.go b/internal/testing/server_token_introspector_mock.go
@@ -9,6 +9,7 @@ package testing
 import (
 	"context"
 	"crypto/sha256"
+	"encoding/base64"
 	"net/http"
 	"net/url"
 
@@ -105,6 +106,7 @@ type GRPCServerTokenIntrospectorMock struct {
 
 	Called                bool
 	LastAuthorizationMeta string
+	LastSessionMeta       string
 	LastRequest           *pb.IntrospectTokenRequest
 }
 
@@ -136,14 +138,27 @@ func (m *GRPCServerTokenIntrospectorMock) IntrospectToken(
 	} else {
 		m.LastAuthorizationMeta = ""
 	}
+	if mdVal := metadata.ValueFromIncomingContext(ctx, "x-session-id"); len(mdVal) != 0 {
+		m.LastSessionMeta = mdVal[0]
+	} else {
+		m.LastSessionMeta = ""
+	}
 	m.LastRequest = req
 
-	if m.LastAuthorizationMeta == "" {
-		return nil, status.Error(codes.Unauthenticated, "Access Token is missing")
+	if m.LastAuthorizationMeta == "" && m.LastSessionMeta == "" {
+		return nil, status.Error(codes.Unauthenticated, "Access Token or Session ID is missing")
 	}
-	if m.LastAuthorizationMeta != "Bearer "+m.accessTokenForIntrospection {
+	if m.LastAuthorizationMeta != "" && m.LastSessionMeta != "" {
+		return nil, status.Error(codes.Unauthenticated, "both Access Token and Session ID cannot be provided")
+	}
+	if m.LastAuthorizationMeta != "" && m.LastAuthorizationMeta != "Bearer "+m.accessTokenForIntrospection {
 		return nil, status.Error(codes.Unauthenticated, "Access Token is invalid")
 	}
+	if m.LastSessionMeta != "" && m.LastSessionMeta != generateSessionID(m.accessTokenForIntrospection) {
+		return nil, status.Error(codes.Unauthenticated, "Session ID is invalid")
+	}
+
+	ctx = metadata.AppendToOutgoingContext(ctx, "x-session-id", generateSessionID(m.accessTokenForIntrospection))
 
 	if result, ok := m.introspectionResults[tokenToKey(req.Token)]; ok {
 		return result, nil
@@ -178,3 +193,8 @@ func (m *GRPCServerTokenIntrospectorMock) ResetCallsInfo() {
 func tokenToKey(token string) [sha256.Size]byte {
 	return sha256.Sum256([]byte(token))
 }
+
+func generateSessionID(token string) string {
+	sha := sha256.Sum256([]byte(token))
+	return base64.StdEncoding.EncodeToString(sha[:32])
+}