session: separate methods for creating vs updating a session

Author: ellemouton

session/interface.go

@@ -123,10 +123,9 @@ func NewSession(label string, typ Type, expiry time.Time, serverAddr string,
 // Store is the interface a persistent storage must implement for storing and
 // retrieving Terminal Connect sessions.
 type Store interface {
-	// StoreSession stores a session in the store. If a session with the
-	// same local public key already exists, the existing record is updated/
-	// overwritten instead.
-	StoreSession(*Session) error
+	// CreateSession adds a new session to the store. If a session with the same
+	// local public key already exists an error is returned.
+	CreateSession(*Session) error
 
 	// GetSession fetches the session with the given key.
 	GetSession(key *btcec.PublicKey) (*Session, error)
@@ -137,4 +136,9 @@ type Store interface {
 	// RevokeSession updates the state of the session with the given local
 	// public key to be revoked.
 	RevokeSession(*btcec.PublicKey) error
+
+	// UpdateSessionRemotePubKey can be used to add the given remote pub key
+	// to the session with the given local pub key.
+	UpdateSessionRemotePubKey(localPubKey,
+		remotePubKey *btcec.PublicKey) error
 }

session/server.go

@@ -33,7 +33,7 @@ func newMailboxSession() *mailboxSession {
 
 func (m *mailboxSession) start(session *Session,
 	serverCreator GRPCServerCreator, authData []byte,
-	onUpdate func(sess *Session) error,
+	onUpdate func(local, remote *btcec.PublicKey) error,
 	onNewStatus func(s mailbox.ServerStatus)) error {
 
 	tlsConfig := &tls.Config{}
@@ -46,8 +46,7 @@ func (m *mailboxSession) start(session *Session,
 	keys := mailbox.NewConnData(
 		ecdh, session.RemotePublicKey, session.PairingSecret[:],
 		authData, func(key *btcec.PublicKey) error {
-			session.RemotePublicKey = key
-			return onUpdate(session)
+			return onUpdate(session.LocalPublicKey, key)
 		}, nil,
 	)
 
@@ -105,7 +104,7 @@ func NewServer(serverCreator GRPCServerCreator) *Server {
 }
 
 func (s *Server) StartSession(session *Session, authData []byte,
-	onUpdate func(sess *Session) error,
+	onUpdate func(local, remote *btcec.PublicKey) error,
 	onNewStatus func(s mailbox.ServerStatus)) (chan struct{}, error) {
 
 	s.activeSessionsMtx.Lock()

session/store.go

@@ -3,6 +3,7 @@ package session
 import (
 	"bytes"
 	"errors"
+	"fmt"
 	"time"
 
 	"github.com/btcsuite/btcd/btcec/v2"
@@ -25,12 +26,11 @@ func getSessionKey(session *Session) []byte {
 	return session.LocalPublicKey.SerializeCompressed()
 }
 
-// StoreSession stores a session in the store. If a session with the
-// same local public key already exists, the existing record is updated/
-// overwritten instead.
+// CreateSession adds a new session to the store. If a session with the same
+// local public key already exists an error is returned.
 //
 // NOTE: this is part of the Store interface.
-func (db *DB) StoreSession(session *Session) error {
+func (db *DB) CreateSession(session *Session) error {
 	var buf bytes.Buffer
 	if err := SerializeSession(&buf, session); err != nil {
 		return err
@@ -43,10 +43,55 @@ func (db *DB) StoreSession(session *Session) error {
 			return err
 		}
 
+		if len(sessionBucket.Get(sessionKey)) != 0 {
+			return fmt.Errorf("session with local public "+
+				"key(%x) already exists",
+				session.LocalPublicKey.SerializeCompressed())
+		}
+
 		return sessionBucket.Put(sessionKey, buf.Bytes())
 	})
 }
 
+// UpdateSessionRemotePubKey can be used to add the given remote pub key
+// to the session with the given local pub key.
+//
+// NOTE: this is part of the Store interface.
+func (db *DB) UpdateSessionRemotePubKey(localPubKey,
+	remotePubKey *btcec.PublicKey) error {
+
+	key := localPubKey.SerializeCompressed()
+
+	return db.Update(func(tx *bbolt.Tx) error {
+		sessionBucket, err := getBucket(tx, sessionBucketKey)
+		if err != nil {
+			return err
+		}
+
+		serialisedSession := sessionBucket.Get(key)
+
+		if len(serialisedSession) == 0 {
+			return ErrSessionNotFound
+		}
+
+		session, err := DeserializeSession(
+			bytes.NewReader(serialisedSession),
+		)
+		if err != nil {
+			return err
+		}
+
+		session.RemotePublicKey = remotePubKey
+
+		var buf bytes.Buffer
+		if err := SerializeSession(&buf, session); err != nil {
+			return err
+		}
+
+		return sessionBucket.Put(key, buf.Bytes())
+	})
+}
+
 // GetSession fetches the session with the given key.
 //
 // NOTE: this is part of the Store interface.
@@ -122,7 +167,7 @@ func (db *DB) ListSessions(filterFn func(s *Session) bool) ([]*Session, error) {
 // NOTE: this is part of the Store interface.
 func (db *DB) RevokeSession(key *btcec.PublicKey) error {
 	var session *Session
-	err := db.View(func(tx *bbolt.Tx) error {
+	return db.Update(func(tx *bbolt.Tx) error {
 		sessionBucket, err := getBucket(tx, sessionBucketKey)
 		if err != nil {
 			return err
@@ -134,14 +179,18 @@ func (db *DB) RevokeSession(key *btcec.PublicKey) error {
 		}
 
 		session, err = DeserializeSession(bytes.NewReader(sessionBytes))
-		return err
-	})
-	if err != nil {
-		return err
-	}
+		if err != nil {
+			return err
+		}
 
-	session.State = StateRevoked
-	session.RevokedAt = time.Now()
+		session.State = StateRevoked
+		session.RevokedAt = time.Now()
 
-	return db.StoreSession(session)
+		var buf bytes.Buffer
+		if err := SerializeSession(&buf, session); err != nil {
+			return err
+		}
+
+		return sessionBucket.Put(key.SerializeCompressed(), buf.Bytes())
+	})
 }

session/store_test.go

@@ -0,0 +1,80 @@
+package session
+
+import (
+	"testing"
+	"time"
+
+	"github.com/btcsuite/btcd/btcec/v2"
+	"github.com/stretchr/testify/require"
+)
+
+// TestBasicSessionStore tests the basic getters and setters of the session
+// store.
+func TestBasicSessionStore(t *testing.T) {
+	// Set up a new DB.
+	db, err := NewDB(t.TempDir(), "test.db")
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = db.Close()
+	})
+
+	// Create a few sessions.
+	s1 := newSession(t, "session 1")
+	s2 := newSession(t, "session 2")
+	s3 := newSession(t, "session 3")
+
+	// Persist session 1.
+	require.NoError(t, db.CreateSession(s1))
+
+	// Trying to persist session 1 again should fail.
+	require.ErrorContains(t, db.CreateSession(s1), "already exists")
+
+	// Persist a few more sessions.
+	require.NoError(t, db.CreateSession(s2))
+	require.NoError(t, db.CreateSession(s3))
+
+	// Ensure that we can retrieve each session.
+	for _, s := range []*Session{s1, s2, s3} {
+		session, err := db.GetSession(s.LocalPublicKey)
+		require.NoError(t, err)
+		require.Equal(t, s.Label, session.Label)
+	}
+
+	// Fetch session 1 and assert that it currently has no remote pub key.
+	session1, err := db.GetSession(s1.LocalPublicKey)
+	require.NoError(t, err)
+	require.Nil(t, session1.RemotePublicKey)
+
+	// Use the update method to add a remote key.
+	remotePriv, err := btcec.NewPrivateKey()
+	require.NoError(t, err)
+	remotePub := remotePriv.PubKey()
+
+	err = db.UpdateSessionRemotePubKey(session1.LocalPublicKey, remotePub)
+	require.NoError(t, err)
+
+	// Assert that the session now does have the remote pub key.
+	session1, err = db.GetSession(s1.LocalPublicKey)
+	require.NoError(t, err)
+	require.True(t, remotePub.IsEqual(session1.RemotePublicKey))
+
+	// Check that the session's state is currently StateCreated.
+	require.Equal(t, session1.State, StateCreated)
+
+	// Now revoke the session and assert that the state is revoked.
+	require.NoError(t, db.RevokeSession(s1.LocalPublicKey))
+	session1, err = db.GetSession(s1.LocalPublicKey)
+	require.NoError(t, err)
+	require.Equal(t, session1.State, StateRevoked)
+}
+
+func newSession(t *testing.T, label string) *Session {
+	session, err := NewSession(
+		label, TypeMacaroonAdmin,
+		time.Date(99999, 1, 1, 0, 0, 0, 0, time.UTC),
+		"foo.bar.baz:1234", true, nil, nil, nil, true,
+	)
+	require.NoError(t, err)
+
+	return session
+}

session_rpcserver.go

@@ -323,7 +323,7 @@ func (s *sessionRpcServer) AddSession(_ context.Context,
 		return nil, fmt.Errorf("error creating new session: %v", err)
 	}
 
-	if err := s.db.StoreSession(sess); err != nil {
+	if err := s.db.CreateSession(sess); err != nil {
 		return nil, fmt.Errorf("error storing session: %v", err)
 	}
 
@@ -477,7 +477,7 @@ func (s *sessionRpcServer) resumeSession(sess *session.Session) error {
 
 	authData := []byte(fmt.Sprintf("%s: %s", HeaderMacaroon, mac))
 	sessionClosedSub, err := s.sessionServer.StartSession(
-		sess, authData, s.db.StoreSession, onNewStatus,
+		sess, authData, s.db.UpdateSessionRemotePubKey, onNewStatus,
 	)
 	if err != nil {
 		return err
@@ -1015,7 +1015,7 @@ func (s *sessionRpcServer) AddAutopilotSession(ctx context.Context,
 	// We only persist this session if we successfully retrieved the
 	// autopilot's static key.
 	sess.RemotePublicKey = remoteKey
-	if err := s.db.StoreSession(sess); err != nil {
+	if err := s.db.CreateSession(sess); err != nil {
 		return nil, fmt.Errorf("error storing session: %v", err)
 	}