db/mongodb: allow implying many params

1 files changed, 31 insertions, 3 deletions

diff --git a/db/mongodb/client.go b/db/mongodb/client.go
index c423d99..1f3815d 100644
--- a/db/mongodb/client.go
+++ b/db/mongodb/client.go
@@ -3,7 +3,9 @@ package mongodb
 import (
         "context"
         "fmt"
+        "net"
         "net/url"
+        "strings"
 
         "go.mongodb.org/mongo-driver/bson"
         "go.mongodb.org/mongo-driver/mongo"
@@ -24,20 +26,46 @@ type Mongo struct {
         db     *mongo.Database
 }
 
-func Connect(ctx context.Context, uri, materialSet string, vc vault.VaultClient) (*Mongo, error) {
+func Connect(ctx context.Context, uri string, vc vault.VaultClient) (*Mongo, error) {
         db := &Mongo{}
 
-        cred, err := vc.DbCredential(ctx, materialSet)
+        // Prefix uri with mongodb:// unless it already includes one of the
+        // standard prefixes (only these two are valid). Otherwise if scheme is
+        // omitted then url parsing will fail to capture the username for Vault
+        // lookup.
+        if !strings.HasPrefix(uri, "mongodb://") && !strings.HasPrefix(uri, "mongodb+srv://") {
+                uri = "mongodb://" + uri
+        }
+
+        u, err := url.Parse(uri)
         if err != nil {
                 return nil, err
         }
 
-        u, err := url.Parse(uri)
+        // The username provided by the user (there should be no
+        // password) will be a reference to a vault material with the
+        // prefix database/creds/. This needs to be replaced with the real
+        // username/password pair fetched from Vault before attempting to
+        // connect.
+        cred, err := vc.DbCredential(ctx, u.User.Username())
         if err != nil {
                 return nil, err
         }
         u.User = url.UserPassword(cred.Username, cred.Password)
 
+        // User may imply the default port
+        if u.Port() == "" {
+                u.Host = net.JoinHostPort(u.Host, "27017")
+        }
+
+        // Users should generally authenticate against the admin collection so
+        // they should only specify it if they need to override that.
+        if u.Query().Get("authSource") == "" {
+                pq := u.Query()
+                pq.Add("authSource", "admin")
+                u.RawQuery = pq.Encode()
+        }
+
         cs, err := connstring.ParseAndValidate(u.String())
         if err != nil {
                 return nil, err