Implement basic dump/restore functionality

server/keystore.go

 package server
 
 import (
+	"compress/gzip"
 	"context"
 	"errors"
+	"io"
 	"io/ioutil"
 	"log"
 	"strings"
@@ -125,6 +127,28 @@ func (s *KeyStore) expireLoop() {
 	}
 }
 
+// Dump the keystore in-memory contents.
+func (s *KeyStore) Dump(w io.Writer) error {
+	gzw := gzip.NewWriter(w)
+	if err := s.userKeys.dump(gzw); err != nil {
+		return err
+	}
+	if err := gzw.Flush(); err != nil {
+		return err
+	}
+	return gzw.Close()
+}
+
+// Load a keystore data dump.
+func (s *KeyStore) Load(r io.Reader) error {
+	gzr, err := gzip.NewReader(r)
+	if err != nil {
+		return err
+	}
+	defer gzr.Close()
+	return s.userKeys.load(gzr)
+}
+
 // Open the user's key store with the given password. If successful,
 // the unencrypted user key will be stored for at most ttlSeconds, or
 // until Close is called with the same session ID.
@@ -195,7 +219,7 @@ func (s *KeyStore) Get(username, ssoTicket string) ([]byte, error) {
 	if !ok {
 		return nil, errNoKeys
 	}
-	return u.pkey, nil
+	return u.PKey, nil
 }
 
 // Close the user's key store and wipe the associated unencrypted key

server/keystore_test.go

@@ -283,3 +283,40 @@ func TestKeystore_Expire(t *testing.T) {
 		t.Fatal("keystore.Get():", err)
 	}
 }
+
+func TestKeystore_DumpAndReload(t *testing.T) {
+	_, keystore, cleanup := newTestContext(t)
+	defer cleanup()
+
+	// Decrypt the private key with the right password.
+	err := keystore.Open(context.Background(), "testuser", string(pw), "session", 60)
+	if err != nil {
+		t.Fatal("keystore.Open():", err)
+	}
+
+	// Now dump the keystore state, and reload it in a new one.
+	var buf bytes.Buffer
+	if err := keystore.Dump(&buf); err != nil {
+		t.Fatal("keystore.Dump():", err)
+	}
+
+	c, keystore2, cleanup2 := newTestContext(t)
+	defer cleanup2()
+
+	if err := keystore2.Load(&buf); err != nil {
+		t.Fatal("keystore2.Load():", err)
+	}
+
+	// Sign a valid SSO ticket and use it to obtain the private
+	// key we just stored.
+	ssoTicket := c.sign("testuser", "keystore/", "domain")
+	result, err := keystore2.Get("testuser", ssoTicket)
+	if err != nil {
+		t.Fatal("keystore2.Get():", err)
+	}
+
+	expectedPEM, _ := privKey.PEM()
+	if !bytes.Equal(result, expectedPEM) {
+		t.Fatalf("keystore2.Get() returned bad key: got %v, expected %v", result, expectedPEM)
+	}
+}

server/map.go

 package server
 
 import (
+	"bytes"
 	"container/heap"
+	"encoding/gob"
+	"io"
 	"time"
 )
 
 type userKey struct {
-	pkey []byte
+	PKey []byte
 }
 
 func newUserKey(pem []byte) *userKey {
-	return &userKey{pkey: pem}
+	return &userKey{PKey: pem}
 }
 
 type userSession struct {
-	id       string
-	username string
-	expiry   time.Time
-	index    int
+	ID       string
+	Username string
+	Expiry   time.Time
+	Index    int
 }
 
 // Priority queue of userSession objects, kept ordered by expiration
@@ -30,18 +33,18 @@ func (pq sessionPQ) Len() int {
 
 func (pq sessionPQ) Swap(i, j int) {
 	pq[i], pq[j] = pq[j], pq[i]
-	pq[i].index = i
-	pq[j].index = j
+	pq[i].Index = i
+	pq[j].Index = j
 }
 
 func (pq sessionPQ) Less(i, j int) bool {
-	return pq[i].expiry.Before(pq[j].expiry)
+	return pq[i].Expiry.Before(pq[j].Expiry)
 }
 
 func (pq *sessionPQ) Push(x interface{}) {
 	n := len(*pq)
 	item := x.(*userSession)
-	item.index = n
+	item.Index = n
 	*pq = append(*pq, item)
 }
 
@@ -67,9 +70,9 @@ func newSessionMap() *sessionMap {
 
 func (m *sessionMap) add(sessionID, username string, ttl time.Duration) {
 	sess := &userSession{
-		id:       sessionID,
-		username: username,
-		expiry:   time.Now().Add(ttl),
+		ID:       sessionID,
+		Username: username,
+		Expiry:   time.Now().Add(ttl),
 	}
 	m.sessions[sessionID] = sess
 	heap.Push(&m.pq, sess)
@@ -81,60 +84,82 @@ func (m *sessionMap) del(sessionID string) *userSession {
 		return nil
 	}
 	delete(m.sessions, sessionID)
-	heap.Remove(&m.pq, sess.index)
-	sess.index = -1
+	heap.Remove(&m.pq, sess.Index)
+	sess.Index = -1
 	return sess
 }
 
 // Peek at the oldest entry and pop it if expired.
 func (m *sessionMap) expireNext(deadline time.Time) *userSession {
-	if len(m.pq) > 0 && m.pq[0].expiry.Before(deadline) {
+	if len(m.pq) > 0 && m.pq[0].Expiry.Before(deadline) {
 		sess := heap.Pop(&m.pq).(*userSession)
-		delete(m.sessions, sess.id)
-		sess.index = -1
+		delete(m.sessions, sess.ID)
+		sess.Index = -1
 		return sess
 	}
 	return nil
 }
 
+func (m *sessionMap) GobEncode() ([]byte, error) {
+	var buf bytes.Buffer
+	if err := gob.NewEncoder(&buf).Encode([]*userSession(m.pq)); err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
+func (m *sessionMap) GobDecode(b []byte) error {
+	var tmp []*userSession
+	if err := gob.NewDecoder(bytes.NewReader(b)).Decode(&tmp); err != nil {
+		return err
+	}
+	sessions := make(map[string]*userSession)
+	for _, sess := range tmp {
+		sessions[sess.ID] = sess
+	}
+	m.pq = sessionPQ(tmp)
+	m.sessions = sessions
+	return nil
+}
+
 // Maintain association between sessions (that may expire) and user
 // keys. Enforces the constraint that keys will be removed once they
 // have no sessions attached.
 type userKeyMap struct {
-	sessions     *sessionMap
-	userKeys     map[string]*userKey
-	userSessions map[string][]string
+	Sessions     *sessionMap
+	UserKeys     map[string]*userKey
+	UserSessions map[string][]string
 }
 
 func newUserKeyMap() *userKeyMap {
 	return &userKeyMap{
-		sessions:     newSessionMap(),
-		userKeys:     make(map[string]*userKey),
-		userSessions: make(map[string][]string),
+		Sessions:     newSessionMap(),
+		UserKeys:     make(map[string]*userKey),
+		UserSessions: make(map[string][]string),
 	}
 }
 
 func (u *userKeyMap) numCachedKeys() int {
-	return len(u.userKeys)
+	return len(u.UserKeys)
 }
 
 func (u *userKeyMap) numSessions() int {
-	return u.sessions.pq.Len()
+	return u.Sessions.pq.Len()
 }
 
 func (u *userKeyMap) get(username string) (*userKey, bool) {
-	k, ok := u.userKeys[username]
+	k, ok := u.UserKeys[username]
 	return k, ok
 }
 
 func (u *userKeyMap) addSessionWithKey(sessionID, username string, key *userKey, ttlSeconds int) {
-	u.sessions.add(sessionID, username, time.Duration(ttlSeconds)*time.Second)
-	u.userKeys[username] = key
-	u.userSessions[username] = append(u.userSessions[username], sessionID)
+	u.Sessions.add(sessionID, username, time.Duration(ttlSeconds)*time.Second)
+	u.UserKeys[username] = key
+	u.UserSessions[username] = append(u.UserSessions[username], sessionID)
 }
 
 func (u *userKeyMap) deleteSession(sessionID string) bool {
-	if sess := u.sessions.del(sessionID); sess != nil {
+	if sess := u.Sessions.del(sessionID); sess != nil {
 		u.cleanupSession(sess)
 		return true
 	}
@@ -143,7 +168,7 @@ func (u *userKeyMap) deleteSession(sessionID string) bool {
 
 func (u *userKeyMap) expire(deadline time.Time) {
 	for {
-		sess := u.sessions.expireNext(deadline)
+		sess := u.Sessions.expireNext(deadline)
 		if sess == nil {
 			return
 		}
@@ -153,21 +178,29 @@ func (u *userKeyMap) expire(deadline time.Time) {
 
 func (u *userKeyMap) cleanupSession(sess *userSession) {
 	var ids []string
-	for _, id := range u.userSessions[sess.username] {
-		if id != sess.id {
+	for _, id := range u.UserSessions[sess.Username] {
+		if id != sess.ID {
 			ids = append(ids, id)
 		}
 	}
 
 	if len(ids) == 0 {
 		// No more sessions for this user, delete key.
-		k := u.userKeys[sess.username]
-		delete(u.userKeys, sess.username)
-		wipeBytes(k.pkey)
-		delete(u.userSessions, sess.username)
+		k := u.UserKeys[sess.Username]
+		delete(u.UserKeys, sess.Username)
+		wipeBytes(k.PKey)
+		delete(u.UserSessions, sess.Username)
 	} else {
-		u.userSessions[sess.username] = ids
+		u.UserSessions[sess.Username] = ids
+	}
 }
+
+func (u *userKeyMap) dump(w io.Writer) error {
+	return gob.NewEncoder(w).Encode(u)
+}
+
+func (u *userKeyMap) load(r io.Reader) error {
+	return gob.NewDecoder(r).Decode(u)
 }
 
 func wipeBytes(b []byte) {
@@ -175,3 +208,13 @@ func wipeBytes(b []byte) {
 		b[i] = 0
 	}
 }
+
+func init() {
+	// Register the types referenced by userKeyMap with the gob
+	// encoder. Use short names so that maps of objects take less
+	// space in the output.
+	gob.RegisterName("k", &userKey{})
+	gob.RegisterName("s", &userSession{})
+	gob.RegisterName("m", &sessionMap{})
+	gob.RegisterName("ukm", &userKeyMap{})
+}