From 3953c2166e538ec2ab5531434f019c0efec8ac33 Mon Sep 17 00:00:00 2001
From: Denis Kirillov <denis@nspcc.ru>
Date: Mon, 30 May 2022 11:40:28 +0300
Subject: [PATCH] [#255] Create session token using arbitrary key

Signed-off-by: Denis Kirillov <denis@nspcc.ru>
---
 client/session.go | 25 +++++++++++++++++++++++--
 pool/pool.go      | 21 +++++++++++++++------
 2 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/client/session.go b/client/session.go
index 8f41891..26f053a 100644
--- a/client/session.go
+++ b/client/session.go
@@ -2,6 +2,7 @@ package client
 
 import (
 	"context"
+	"crypto/ecdsa"
 
 	"github.com/nspcc-dev/neofs-api-go/v2/refs"
 	rpcapi "github.com/nspcc-dev/neofs-api-go/v2/rpc"
@@ -15,6 +16,9 @@ type PrmSessionCreate struct {
 	prmCommonMeta
 
 	exp uint64
+
+	keySet bool
+	key    ecdsa.PrivateKey
 }
 
 // SetExp sets number of the last NepFS epoch in the lifetime of the session after which it will be expired.
@@ -22,6 +26,13 @@ func (x *PrmSessionCreate) SetExp(exp uint64) {
 	x.exp = exp
 }
 
+// UseKey specifies private key to sign the requests and compute token owner.
+// If key is not provided, then Client default key is used.
+func (x *PrmSessionCreate) UseKey(key ecdsa.PrivateKey) {
+	x.keySet = true
+	x.key = key
+}
+
 // ResSessionCreate groups resulting values of SessionCreate operation.
 type ResSessionCreate struct {
 	statusRes
@@ -72,8 +83,12 @@ func (c *Client) SessionCreate(ctx context.Context, prm PrmSessionCreate) (*ResS
 		panic(panicMsgMissingContext)
 	}
 
+	ownerKey := c.prm.key.PublicKey
+	if prm.keySet {
+		ownerKey = prm.key.PublicKey
+	}
 	var ownerID user.ID
-	user.IDFromKey(&ownerID, c.prm.key.PublicKey)
+	user.IDFromKey(&ownerID, ownerKey)
 
 	var ownerIDV2 refs.OwnerID
 	ownerID.WriteToV2(&ownerIDV2)
@@ -95,7 +110,13 @@ func (c *Client) SessionCreate(ctx context.Context, prm PrmSessionCreate) (*ResS
 		res ResSessionCreate
 	)
 
-	c.initCallContext(&cc)
+	if prm.keySet {
+		c.initCallContextWithoutKey(&cc)
+		cc.key = prm.key
+	} else {
+		c.initCallContext(&cc)
+	}
+
 	cc.meta = prm.prmCommonMeta
 	cc.req = &req
 	cc.statusRes = &res
diff --git a/pool/pool.go b/pool/pool.go
index 10c628d..8856b1d 100644
--- a/pool/pool.go
+++ b/pool/pool.go
@@ -428,6 +428,7 @@ func (c *clientWrapper) objectSearch(ctx context.Context, prm PrmObjectSearch) (
 func (c *clientWrapper) sessionCreate(ctx context.Context, prm prmCreateSession) (*resCreateSession, error) {
 	var cliPrm sdkClient.PrmSessionCreate
 	cliPrm.SetExp(prm.exp)
+	cliPrm.UseKey(prm.key)
 
 	res, err := c.client.SessionCreate(ctx, cliPrm)
 	if err != nil {
@@ -835,13 +836,20 @@ func (x *PrmBalanceGet) SetAccount(id user.ID) {
 // prmEndpointInfo groups parameters of sessionCreate operation.
 type prmCreateSession struct {
 	exp uint64
+	key ecdsa.PrivateKey
 }
 
-// SetExp sets number of the last NeoFS epoch in the lifetime of the session after which it will be expired.
-func (x *prmCreateSession) SetExp(exp uint64) {
+// setExp sets number of the last NeoFS epoch in the lifetime of the session after which it will be expired.
+func (x *prmCreateSession) setExp(exp uint64) {
 	x.exp = exp
 }
 
+// useKey specifies owner private key for session token.
+// If key is not provided, then Pool default key is used.
+func (x *prmCreateSession) useKey(key ecdsa.PrivateKey) {
+	x.key = key
+}
+
 // prmEndpointInfo groups parameters of endpointInfo operation.
 type prmEndpointInfo struct{}
 
@@ -957,7 +965,7 @@ func (p *Pool) Dial(ctx context.Context) error {
 			}
 			var healthy bool
 			var st session.Object
-			err = initSessionForDuration(ctx, &st, c, p.rebalanceParams.sessionExpirationDuration)
+			err = initSessionForDuration(ctx, &st, c, p.rebalanceParams.sessionExpirationDuration, *p.key)
 			if err != nil && p.logger != nil {
 				p.logger.Warn("failed to create neofs session token for client",
 					zap.String("Address", addr),
@@ -1199,7 +1207,7 @@ func (p *Pool) checkSessionTokenErr(err error, address string) bool {
 	return false
 }
 
-func initSessionForDuration(ctx context.Context, dst *session.Object, c client, dur uint64) error {
+func initSessionForDuration(ctx context.Context, dst *session.Object, c client, dur uint64, ownerKey ecdsa.PrivateKey) error {
 	ni, err := c.networkInfo(ctx, prmNetworkInfo{})
 	if err != nil {
 		return err
@@ -1214,7 +1222,8 @@ func initSessionForDuration(ctx context.Context, dst *session.Object, c client,
 		exp = epoch + dur
 	}
 	var prm prmCreateSession
-	prm.SetExp(exp)
+	prm.setExp(exp)
+	prm.useKey(ownerKey)
 
 	res, err := c.sessionCreate(ctx, prm)
 	if err != nil {
@@ -1302,7 +1311,7 @@ func (p *Pool) openDefaultSession(ctx *callContext) error {
 	tok, ok := p.cache.Get(cacheKey)
 	if !ok {
 		// init new session
-		err := initSessionForDuration(ctx, &tok, ctx.client, p.stokenDuration)
+		err := initSessionForDuration(ctx, &tok, ctx.client, p.stokenDuration, *ctx.key)
 		if err != nil {
 			return fmt.Errorf("session API client: %w", err)
 		}