From b556d54f16acd1de144379329a49c1c1e1570623 Mon Sep 17 00:00:00 2001
From: alexvanin <alexey@nspcc.ru>
Date: Fri, 17 Jan 2020 12:39:00 +0300
Subject: [PATCH] session: Modify verify routine to check owner id connection
 with public keys

---
 session/store_test.go |  6 ++++--
 session/types.go      | 23 +++++++++++++++++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/session/store_test.go b/session/store_test.go
index 1a9e977..66f99d4 100644
--- a/session/store_test.go
+++ b/session/store_test.go
@@ -31,20 +31,22 @@ func newTestClient(t *testing.T) *testClient {
 
 func signToken(t *testing.T, token *PToken, c *testClient) {
 	require.NotNil(t, token)
+	token.SetPublicKeys(&c.PublicKey)
 
 	signH, err := c.Sign(token.Header.PublicKey)
 	require.NoError(t, err)
 	require.NotNil(t, signH)
 
 	// data is not yet signed
-	require.False(t, token.Verify(&c.PublicKey))
+	keys := UnmarshalPublicKeys(&token.Token)
+	require.False(t, token.Verify(keys...))
 
 	signT, err := c.Sign(token.verificationData())
 	require.NoError(t, err)
 	require.NotNil(t, signT)
 
 	token.AddSignatures(signH, signT)
-	require.True(t, token.Verify(&c.PublicKey))
+	require.True(t, token.Verify(keys...))
 }
 
 func TestTokenStore(t *testing.T) {
diff --git a/session/types.go b/session/types.go
index aff7cca..08abec8 100644
--- a/session/types.go
+++ b/session/types.go
@@ -6,6 +6,7 @@ import (
 	"sync"
 
 	crypto "github.com/nspcc-dev/neofs-crypto"
+	"github.com/nspcc-dev/neofs-proto/chain"
 	"github.com/nspcc-dev/neofs-proto/internal"
 	"github.com/nspcc-dev/neofs-proto/refs"
 	"github.com/pkg/errors"
@@ -111,11 +112,24 @@ func (m *Token) Sign(key *ecdsa.PrivateKey) error {
 	return nil
 }
 
+// SetPublicKeys sets owner's public keys to the token
+func (m *Token) SetPublicKeys(keys... *ecdsa.PublicKey) {
+	m.PublicKeys = m.PublicKeys[:0]
+	for i := range keys {
+		m.PublicKeys = append(m.PublicKeys, crypto.MarshalPublicKey(keys[i]))
+	}
+}
+
 // Verify checks if token is correct and signed.
 func (m *Token) Verify(keys ...*ecdsa.PublicKey) bool {
 	if m.FirstEpoch > m.LastEpoch {
 		return false
 	}
+	ownerFromKeys := chain.KeysToAddress(keys...)
+	if m.OwnerID.String() != ownerFromKeys {
+		return false
+	}
+
 	for i := range keys {
 		if m.Header.Verify(keys[i]) && crypto.Verify(keys[i], m.verificationData(), m.Signature) == nil {
 			return true
@@ -156,3 +170,12 @@ func (m *VerificationHeader) Verify(keys ...*ecdsa.PublicKey) bool {
 	}
 	return false
 }
+
+// UnmarshalPublicKeys returns unmarshal public keys from the token
+func UnmarshalPublicKeys(t *Token) []*ecdsa.PublicKey {
+	r := make([]*ecdsa.PublicKey, 0, len(t.PublicKeys))
+	for i := range t.PublicKeys {
+		r = append(r, crypto.UnmarshalPublicKey(t.PublicKeys[i]))
+	}
+	return r
+}