From f20da8ed234f44979b430a1f74ea42fa84d9cdc3 Mon Sep 17 00:00:00 2001
From: = <torjus@usit.uio.no>
Date: Mon, 6 Dec 2021 17:28:48 +0100
Subject: [PATCH] Add revocation to certificate stores

---
 store/bolt.go       | 27 +++++++++++++++++++++++++++
 store/memory.go     | 43 +++++++++++++++++++++++++++++++------------
 store/store.go      |  2 ++
 store/store_test.go | 22 +++++++++++++++++++++-
 4 files changed, 81 insertions(+), 13 deletions(-)

diff --git a/store/bolt.go b/store/bolt.go
index e1f1b75..5c33ef5 100644
--- a/store/bolt.go
+++ b/store/bolt.go
@@ -21,6 +21,7 @@ var bktKey = []byte("files")
 var bktKeyCerts = []byte("certs")
 var bktKeyKeys = []byte("keys")
 var bktKeyUsers = []byte("users")
+var bktKeyRevoked = []byte("revoked")
 
 func NewBoltStore(path string) (*BoltStore, error) {
 	s := &BoltStore{}
@@ -42,6 +43,9 @@ func NewBoltStore(path string) (*BoltStore, error) {
 		if _, err := t.CreateBucketIfNotExists(bktKeyUsers); err != nil {
 			return err
 		}
+		if _, err := t.CreateBucketIfNotExists(bktKeyRevoked); err != nil {
+			return err
+		}
 		return nil
 	})
 	if err != nil {
@@ -202,6 +206,29 @@ func (s *BoltStore) ListCertificates() ([]string, error) {
 	return ids, nil
 }
 
+func (s *BoltStore) Revoke(serial string) error {
+	return s.db.Update(func(tx *bolt.Tx) error {
+		bkt := tx.Bucket(bktKeyRevoked)
+		return bkt.Put([]byte(serial), []byte{'r'})
+	})
+}
+
+func (s *BoltStore) IsRevoked(serial string) (bool, error) {
+	var revoked bool
+	err := s.db.View(func(tx *bolt.Tx) error {
+		bkt := tx.Bucket(bktKeyRevoked)
+		status := bkt.Get([]byte(serial))
+		if status != nil {
+			revoked = true
+		}
+		return nil
+	})
+	if err != nil {
+		return false, err
+	}
+	return revoked, nil
+}
+
 var _ UserStore = &BoltStore{}
 
 func (s *BoltStore) StoreUser(user *pb.User) error {
diff --git a/store/memory.go b/store/memory.go
index 4d24976..eeaa617 100644
--- a/store/memory.go
+++ b/store/memory.go
@@ -10,22 +10,25 @@ import (
 )
 
 type MemoryStore struct {
-	filesLock sync.RWMutex
-	files     map[string]*pb.File
-	certLock  sync.RWMutex
-	certs     map[string][]byte
-	keyLock   sync.RWMutex
-	keys      map[string][]byte
-	usersLock sync.RWMutex
-	users     map[string]*pb.User
+	filesLock    sync.RWMutex
+	files        map[string]*pb.File
+	certLock     sync.RWMutex
+	certs        map[string][]byte
+	keyLock      sync.RWMutex
+	keys         map[string][]byte
+	usersLock    sync.RWMutex
+	users        map[string]*pb.User
+	revokedCerts map[string]struct{}
+	revokedLock  sync.RWMutex
 }
 
 func NewMemoryStore() *MemoryStore {
 	return &MemoryStore{
-		files: make(map[string]*pb.File),
-		certs: make(map[string][]byte),
-		keys:  make(map[string][]byte),
-		users: make(map[string]*pb.User),
+		files:        make(map[string]*pb.File),
+		certs:        make(map[string][]byte),
+		keys:         make(map[string][]byte),
+		users:        make(map[string]*pb.User),
+		revokedCerts: make(map[string]struct{}),
 	}
 }
 
@@ -147,6 +150,22 @@ func (s *MemoryStore) ListCertificates() ([]string, error) {
 	return certIDs, nil
 }
 
+func (s *MemoryStore) Revoke(serial string) error {
+	s.revokedLock.Lock()
+	defer s.revokedLock.Unlock()
+
+	s.revokedCerts[serial] = struct{}{}
+	return nil
+}
+
+func (s *MemoryStore) IsRevoked(serial string) (bool, error) {
+	s.revokedLock.RLock()
+	defer s.revokedLock.RUnlock()
+
+	_, revoked := s.revokedCerts[serial]
+	return revoked, nil
+}
+
 ///////////////
 // UserStore //
 ///////////////
diff --git a/store/store.go b/store/store.go
index e5687b5..d2d0873 100644
--- a/store/store.go
+++ b/store/store.go
@@ -23,6 +23,8 @@ type CertificateStore interface {
 	GetKey(id string) (*ecdsa.PrivateKey, error)
 	StoreKey(id string, key *ecdsa.PrivateKey) error
 	ListCertificates() ([]string, error)
+	Revoke(serial string) error
+	IsRevoked(serial string) (bool, error)
 }
 
 type UserStore interface {
diff --git a/store/store_test.go b/store/store_test.go
index 6c3e99a..2b5c199 100644
--- a/store/store_test.go
+++ b/store/store_test.go
@@ -79,7 +79,7 @@ func doCertificateStoreTest(s store.CertificateStore, t *testing.T) {
 
 		// Create cert and key
 		unsigned := &x509.Certificate{
-			SerialNumber: big.NewInt(time.Now().Unix()),
+			SerialNumber: big.NewInt(time.Now().UnixMilli()),
 			Subject: pkix.Name{
 				Organization: []string{"ezshare"},
 				Country:      []string{"No"},
@@ -142,6 +142,26 @@ func doCertificateStoreTest(s store.CertificateStore, t *testing.T) {
 		if !retrievedKey.Equal(privateKey) {
 			t.Errorf("Retrieved key does not match stored.")
 		}
+
+		// Revoke
+		isRevoked, err := s.IsRevoked(cert.SerialNumber.String())
+		if err != nil {
+			t.Fatalf("Error checking if certificate is revoked: %s", err)
+		}
+		if isRevoked {
+			t.Fatalf("Unrevoked certificate is revoked")
+		}
+		if err := s.Revoke(cert.SerialNumber.String()); err != nil {
+			t.Fatalf("Error revoking certificate: %s", err)
+		}
+		isRevoked, err = s.IsRevoked(cert.SerialNumber.String())
+		if err != nil {
+			t.Fatalf("Error checking if certificate is revoked: %s", err)
+		}
+		if !isRevoked {
+			t.Fatalf("Revoked certificate is not revoked")
+		}
+
 	})
 }