From 0259a06783bb11b0192f5ebbdb55a343aa8c937c Mon Sep 17 00:00:00 2001
From: Leonard Lyubich <leonard@nspcc.ru>
Date: Thu, 17 Sep 2020 18:04:04 +0300
Subject: [PATCH] [#150] sdk/token: Add owner ID

Signed-off-by: Leonard Lyubich <leonard@nspcc.ru>
---
 pkg/client/object.go      |  8 ++--
 pkg/client/session.go     |  8 ++--
 pkg/token/session.go      | 80 +++++++++++++++++++++++++++------------
 pkg/token/session_test.go | 42 ++++++++++++++++++++
 4 files changed, 104 insertions(+), 34 deletions(-)
 create mode 100644 pkg/token/session_test.go

diff --git a/pkg/client/object.go b/pkg/client/object.go
index f7f95a0..8bcedb2 100644
--- a/pkg/client/object.go
+++ b/pkg/client/object.go
@@ -1034,6 +1034,8 @@ func (c Client) attachV2SessionToken(opts callOptions, hdr *v2session.RequestMet
 		return nil
 	}
 
+	token := opts.session.ToV2()
+
 	opCtx := new(v2session.ObjectSessionContext)
 	opCtx.SetAddress(info.addr)
 	opCtx.SetVerb(info.verb)
@@ -1043,15 +1045,11 @@ func (c Client) attachV2SessionToken(opts callOptions, hdr *v2session.RequestMet
 	lt.SetNbf(info.nbf)
 	lt.SetExp(info.exp)
 
-	body := new(v2session.SessionTokenBody)
-	body.SetID(opts.session.ID())
+	body := token.GetBody()
 	body.SetSessionKey(opts.session.SessionKey())
 	body.SetContext(opCtx)
 	body.SetLifetime(lt)
 
-	token := new(v2session.SessionToken)
-	token.SetBody(body)
-
 	signWrapper := signature.StableMarshalerWrapper{SM: token.GetBody()}
 	err := signer.SignDataWithHandler(c.key, signWrapper, func(key []byte, sig []byte) {
 		sessionTokenSignature := new(v2refs.Signature)
diff --git a/pkg/client/session.go b/pkg/client/session.go
index fee9e7f..6c912a0 100644
--- a/pkg/client/session.go
+++ b/pkg/client/session.go
@@ -70,10 +70,10 @@ func (c Client) createSessionV2(ctx context.Context, expiration uint64, opts ...
 			return nil, errors.New("malformed response body")
 		}
 
-		sessionToken, err := token.CreateSessionToken(body.GetID(), body.GetSessionKey())
-		if err != nil {
-			return nil, errors.Wrap(err, "malformed response body")
-		}
+		sessionToken := token.NewSessionToken()
+		sessionToken.SetID(body.GetID())
+		sessionToken.SetSessionKey(body.GetSessionKey())
+		sessionToken.SetOwnerID(ownerID)
 
 		return sessionToken, nil
 	default:
diff --git a/pkg/token/session.go b/pkg/token/session.go
index 450a9a7..6ff3753 100644
--- a/pkg/token/session.go
+++ b/pkg/token/session.go
@@ -1,40 +1,70 @@
 package token
 
 import (
-	"github.com/google/uuid"
+	"github.com/nspcc-dev/neofs-api-go/pkg/owner"
+	"github.com/nspcc-dev/neofs-api-go/v2/session"
 )
 
-type SessionToken struct {
-	id     uuid.UUID
-	pubKey []byte
+type SessionToken session.SessionToken
+
+func NewSessionTokenFromV2(tV2 *session.SessionToken) *SessionToken {
+	return (*SessionToken)(tV2)
 }
 
-func CreateSessionToken(id, pub []byte) (*SessionToken, error) {
-	var tokenID uuid.UUID
+func NewSessionToken() *SessionToken {
+	return NewSessionTokenFromV2(new(session.SessionToken))
+}
 
-	err := tokenID.UnmarshalBinary(id)
-	if err != nil {
-		return nil, err
+func (t *SessionToken) ToV2() *session.SessionToken {
+	return (*session.SessionToken)(t)
+}
+
+func (t *SessionToken) setBodyField(setter func(*session.SessionTokenBody)) {
+	token := (*session.SessionToken)(t)
+	body := token.GetBody()
+
+	if body == nil {
+		body = new(session.SessionTokenBody)
+		token.SetBody(body)
 	}
 
-	key := make([]byte, len(pub))
-	copy(key[:], pub)
-
-	return &SessionToken{
-		id:     tokenID,
-		pubKey: key,
-	}, nil
+	setter(body)
 }
 
-func (s SessionToken) SessionKey() []byte {
-	return s.pubKey
+func (t *SessionToken) ID() []byte {
+	return (*session.SessionToken)(t).
+		GetBody().
+		GetID()
 }
 
-func (s SessionToken) ID() []byte {
-	data, err := s.id.MarshalBinary()
-	if err != nil {
-		panic(err) // must never panic
-	}
-
-	return data
+func (t *SessionToken) SetID(v []byte) {
+	t.setBodyField(func(body *session.SessionTokenBody) {
+		body.SetID(v)
+	})
+}
+
+func (t *SessionToken) OwnerID() *owner.ID {
+	return owner.NewIDFromV2(
+		(*session.SessionToken)(t).
+			GetBody().
+			GetOwnerID(),
+	)
+}
+
+func (t *SessionToken) SetOwnerID(v *owner.ID) {
+	t.setBodyField(func(body *session.SessionTokenBody) {
+		body.SetOwnerID(v.ToV2())
+	})
+}
+
+func (t *SessionToken) SessionKey() []byte {
+	return (*session.SessionToken)(t).
+		GetBody().
+		GetSessionKey()
+}
+
+func (t *SessionToken) SetSessionKey(v []byte) {
+	t.setBodyField(func(body *session.SessionTokenBody) {
+		body.SetSessionKey(v)
+	})
 }
diff --git a/pkg/token/session_test.go b/pkg/token/session_test.go
new file mode 100644
index 0000000..f155bc9
--- /dev/null
+++ b/pkg/token/session_test.go
@@ -0,0 +1,42 @@
+package token
+
+import (
+	"crypto/rand"
+	"testing"
+
+	"github.com/nspcc-dev/neofs-api-go/pkg/owner"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSessionToken_SetID(t *testing.T) {
+	token := NewSessionToken()
+
+	id := []byte{1, 2, 3}
+	token.SetID(id)
+
+	require.Equal(t, id, token.ID())
+}
+
+func TestSessionToken_SetOwnerID(t *testing.T) {
+	token := NewSessionToken()
+
+	w := new(owner.NEO3Wallet)
+	_, err := rand.Read(w.Bytes())
+	require.NoError(t, err)
+
+	ownerID := owner.NewID()
+	ownerID.SetNeo3Wallet(w)
+
+	token.SetOwnerID(ownerID)
+
+	require.Equal(t, ownerID, token.OwnerID())
+}
+
+func TestSessionToken_SetSessionKey(t *testing.T) {
+	token := NewSessionToken()
+
+	key := []byte{1, 2, 3}
+	token.SetSessionKey(key)
+
+	require.Equal(t, key, token.SessionKey())
+}