session: add session get and put helpers

Author: ellemouton

session/kvdb_store.go

@@ -206,10 +206,6 @@ func (db *BoltStore) NewSession(id ID, localPrivKey *btcec.PrivateKey,
 //
 // NOTE: this is part of the Store interface.
 func (db *BoltStore) CreateSession(session *Session) error {
-	var buf bytes.Buffer
-	if err := SerializeSession(&buf, session); err != nil {
-		return err
-	}
 	sessionKey := getSessionKey(session)
 
 	return db.Update(func(tx *bbolt.Tx) error {
@@ -246,21 +242,7 @@ func (db *BoltStore) CreateSession(session *Session) error {
 			}
 
 			for _, id := range sessionIDs {
-				keyBytes, err := getKeyForID(
-					sessionBucket, id,
-				)
-				if err != nil {
-					return err
-				}
-
-				v := sessionBucket.Get(keyBytes)
-				if len(v) == 0 {
-					return ErrSessionNotFound
-				}
-
-				sess, err := DeserializeSession(
-					bytes.NewReader(v),
-				)
+				sess, err := getSessionByID(sessionBucket, id)
 				if err != nil {
 					return err
 				}
@@ -291,7 +273,7 @@ func (db *BoltStore) CreateSession(session *Session) error {
 			return err
 		}
 
-		return sessionBucket.Put(sessionKey, buf.Bytes())
+		return putSession(sessionBucket, session)
 	})
 }
 
@@ -325,12 +307,7 @@ func (db *BoltStore) UpdateSessionRemotePubKey(localPubKey,
 
 		session.RemotePublicKey = remotePubKey
 
-		var buf bytes.Buffer
-		if err := SerializeSession(&buf, session); err != nil {
-			return err
-		}
-
-		return sessionBucket.Put(key, buf.Bytes())
+		return putSession(sessionBucket, session)
 	})
 }
 
@@ -441,7 +418,7 @@ func (db *BoltStore) listSessions(filterFn func(s *Session) bool) ([]*Session, e
 // state.
 //
 // NOTE: this is part of the Store interface.
-func (db *DB) DeleteReservedSessions() error {
+func (db *BoltStore) DeleteReservedSessions() error {
 	return db.Update(func(tx *bbolt.Tx) error {
 		sessionBucket, err := getBucket(tx, sessionBucketKey)
 		if err != nil {
@@ -555,12 +532,7 @@ func (db *BoltStore) RevokeSession(key *btcec.PublicKey) error {
 		session.State = StateRevoked
 		session.RevokedAt = db.clock.Now().UTC()
 
-		var buf bytes.Buffer
-		if err := SerializeSession(&buf, session); err != nil {
-			return err
-		}
-
-		return sessionBucket.Put(key.SerializeCompressed(), buf.Bytes())
+		return putSession(sessionBucket, session)
 	})
 }
 
@@ -575,22 +547,9 @@ func (db *BoltStore) GetSessionByID(id ID) (*Session, error) {
 			return err
 		}
 
-		keyBytes, err := getKeyForID(sessionBucket, id)
-		if err != nil {
-			return err
-		}
-
-		v := sessionBucket.Get(keyBytes)
-		if len(v) == 0 {
-			return ErrSessionNotFound
-		}
+		session, err = getSessionByID(sessionBucket, id)
 
-		session, err = DeserializeSession(bytes.NewReader(v))
-		if err != nil {
-			return err
-		}
-
-		return nil
+		return err
 	})
 	if err != nil {
 		return nil, err
@@ -897,3 +856,26 @@ func addIDToGroupIDPair(sessionBkt *bbolt.Bucket, id, groupID ID) error {
 
 	return sessionIDsBkt.Put(seqNoBytes[:], id[:])
 }
+
+func getSessionByID(bucket *bbolt.Bucket, id ID) (*Session, error) {
+	keyBytes, err := getKeyForID(bucket, id)
+	if err != nil {
+		return nil, err
+	}
+
+	v := bucket.Get(keyBytes)
+	if len(v) == 0 {
+		return nil, ErrSessionNotFound
+	}
+
+	return DeserializeSession(bytes.NewReader(v))
+}
+
+func putSession(bucket *bbolt.Bucket, session *Session) error {
+	var buf bytes.Buffer
+	if err := SerializeSession(&buf, session); err != nil {
+		return err
+	}
+
+	return bucket.Put(getSessionKey(session), buf.Bytes())
+}