Support mongo datastores

4 files changed, 162 insertions, 2 deletions

diff --git a/server/store/config_test.go b/server/store/config_test.go
new file mode 100644
index 0000000..8e283f5
--- /dev/null
+++ b/server/store/config_test.go
@@ -0,0 +1,68 @@
+package store
+
+import (
+	"reflect"
+	"testing"
+	"time"
+
+	mgo "gopkg.in/mgo.v2"
+)
+
+func TestMySQLConfig(t *testing.T) {
+	var tests = []struct {
+		in  string
+		out string
+	}{
+		{"mysql:user:passwd:localhost", "user:passwd@tcp(localhost:3306)/certs?parseTime=true"},
+		{"mysql:user:passwd:localhost:13306", "user:passwd@tcp(localhost:13306)/certs?parseTime=true"},
+		{"mysql:root::localhost", "root@tcp(localhost:3306)/certs?parseTime=true"},
+	}
+	for _, tt := range tests {
+		result := parseMySQLConfig(tt.in)
+		if result != tt.out {
+			t.Errorf("want %s, got %s", tt.out, result)
+		}
+	}
+}
+
+func TestMongoConfig(t *testing.T) {
+	var tests = []struct {
+		in  string
+		out *mgo.DialInfo
+	}{
+		{"mongo:user:passwd:host", &mgo.DialInfo{
+			Username: "user",
+			Password: "passwd",
+			Addrs:    []string{"host"},
+			Database: "certs",
+			Timeout:  5 * time.Second,
+		}},
+		{"mongo:user:passwd:host1,host2", &mgo.DialInfo{
+			Username: "user",
+			Password: "passwd",
+			Addrs:    []string{"host1", "host2"},
+			Database: "certs",
+			Timeout:  5 * time.Second,
+		}},
+		{"mongo:user:passwd:host1:27017,host2:27017", &mgo.DialInfo{
+			Username: "user",
+			Password: "passwd",
+			Addrs:    []string{"host1:27017", "host2:27017"},
+			Database: "certs",
+			Timeout:  5 * time.Second,
+		}},
+		{"mongo:user:passwd:host1,host2:27017", &mgo.DialInfo{
+			Username: "user",
+			Password: "passwd",
+			Addrs:    []string{"host1", "host2:27017"},
+			Database: "certs",
+			Timeout:  5 * time.Second,
+		}},
+	}
+	for _, tt := range tests {
+		result := parseMongoConfig(tt.in)
+		if !reflect.DeepEqual(result, tt.out) {
+			t.Errorf("want:\n%+v\ngot:\n%+v", tt.out, result)
+		}
+	}
+}
diff --git a/server/store/mongo.go b/server/store/mongo.go
new file mode 100644
index 0000000..752d405
--- /dev/null
+++ b/server/store/mongo.go
@@ -0,0 +1,80 @@
+package store
+
+import (
+	"strings"
+	"time"
+
+	"golang.org/x/crypto/ssh"
+
+	mgo "gopkg.in/mgo.v2"
+	"gopkg.in/mgo.v2/bson"
+)
+
+var (
+	certsDB     = "certs"
+	issuedTable = "issued_certs"
+)
+
+type mongoDB struct {
+	conn *mgo.Collection
+}
+
+func parseMongoConfig(config string) *mgo.DialInfo {
+	s := strings.SplitN(config, ":", 4)
+	_, user, passwd, hosts := s[0], s[1], s[2], s[3]
+	d := &mgo.DialInfo{
+		Addrs:    strings.Split(hosts, ","),
+		Username: user,
+		Password: passwd,
+		Database: certsDB,
+		Timeout:  time.Second * 5,
+	}
+	return d
+}
+
+func NewMongoStore(config string) (CertStorer, error) {
+	session, err := mgo.DialWithInfo(parseMongoConfig(config))
+	if err != nil {
+		return nil, err
+	}
+	c := session.DB(certsDB).C(issuedTable)
+	return &mongoDB{
+		conn: c,
+	}, nil
+}
+
+func (m *mongoDB) Get(id string) (*CertRecord, error) {
+	c := &CertRecord{}
+	err := m.conn.Find(bson.M{"keyid": id}).One(c)
+	return c, err
+}
+
+func (m *mongoDB) SetCert(cert *ssh.Certificate) error {
+	r := parseCertificate(cert)
+	return m.SetRecord(r)
+}
+
+func (m *mongoDB) SetRecord(record *CertRecord) error {
+	return m.conn.Insert(record)
+}
+
+func (m *mongoDB) List() ([]*CertRecord, error) {
+	var result []*CertRecord
+	m.conn.Find(nil).All(&result)
+	return result, nil
+}
+
+func (m *mongoDB) Revoke(id string) error {
+	return m.conn.Update(bson.M{"keyid": id}, bson.M{"$set": bson.M{"revoked": true}})
+}
+
+func (m *mongoDB) GetRevoked() ([]*CertRecord, error) {
+	var result []*CertRecord
+	err := m.conn.Find(bson.M{"expires": bson.M{"$gte": time.Now().UTC()}, "revoked": true}).All(&result)
+	return result, err
+}
+
+func (m *mongoDB) Close() error {
+	m.conn.Database.Session.Close()
+	return nil
+}
diff --git a/server/store/mysql.go b/server/store/mysql.go
index a62af6b..7a0b111 100644
--- a/server/store/mysql.go
+++ b/server/store/mysql.go
@@ -22,7 +22,7 @@ type mysqlDB struct {
 	revoked *sql.Stmt
 }
 
-func parseConfig(config string) string {
+func parseMySQLConfig(config string) string {
 	s := strings.Split(config, ":")
 	if len(s) == 4 {
 		s = append(s, "3306")
@@ -41,7 +41,7 @@ func parseConfig(config string) string {
 
 // NewMySQLStore returns a MySQL CertStorer.
 func NewMySQLStore(config string) (CertStorer, error) {
-	conn, err := sql.Open("mysql", parseConfig(config))
+	conn, err := sql.Open("mysql", parseMySQLConfig(config))
 	if err != nil {
 		return nil, fmt.Errorf("mysql: could not get a connection: %v", err)
 	}
diff --git a/server/store/store_test.go b/server/store/store_test.go
index ee80241..bf16fa6 100644
--- a/server/store/store_test.go
+++ b/server/store/store_test.go
@@ -100,3 +100,15 @@ func TestMySQLStore(t *testing.T) {
 	}
 	testStore(t, db)
 }
+
+func TestMongoStore(t *testing.T) {
+	config := os.Getenv("MONGO_TEST_CONFIG")
+	if config == "" {
+		t.Skip("No MONGO_TEST_CONFIG environment variable")
+	}
+	db, err := NewMongoStore(config)
+	if err != nil {
+		t.Error(err)
+	}
+	testStore(t, db)
+}