forked from mongodb/mongo-go-driver
/
mongocryptd.go
137 lines (115 loc) · 3.87 KB
/
mongocryptd.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
// Copyright (C) MongoDB, Inc. 2017-present.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
package mongo
import (
"context"
"os/exec"
"strings"
"time"
"go.mongodb.org/mongo-driver/mongo/options"
"go.mongodb.org/mongo-driver/mongo/readconcern"
"go.mongodb.org/mongo-driver/mongo/readpref"
"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
)
const (
defaultServerSelectionTimeout = 1 * time.Second
defaultURI = "mongodb://localhost:27020"
defaultPath = "mongocryptd"
serverSelectionTimeoutStr = "server selection error"
)
var defaultTimeoutArgs = []string{"--idleShutdownTimeoutSecs=60"}
var databaseOpts = options.Database().SetReadConcern(readconcern.New()).SetReadPreference(readpref.Primary())
type mcryptClient struct {
bypassSpawn bool
client *Client
path string
spawnArgs []string
}
func newMcryptClient(opts map[string]interface{}) (*mcryptClient, error) {
// create mcryptClient instance and spawn process if necessary
mc := &mcryptClient{}
if bypass, ok := opts["mongocryptdBypassSpawn"]; ok {
mc.bypassSpawn = bypass.(bool)
}
if !mc.bypassSpawn {
mc.path, mc.spawnArgs = createSpawnArgs(opts)
if err := mc.spawnProcess(); err != nil {
return nil, err
}
}
// get connection string
uri := defaultURI
if u, ok := opts["mongocryptdURI"]; ok {
uri = u.(string)
}
// create client
client, err := NewClient(options.Client().ApplyURI(uri).SetServerSelectionTimeout(defaultServerSelectionTimeout))
if err != nil {
return nil, err
}
mc.client = client
return mc, nil
}
// markCommand executes the given command on mongocryptd.
func (mc *mcryptClient) markCommand(ctx context.Context, dbName string, cmd bsoncore.Document) (bsoncore.Document, error) {
db := mc.client.Database(dbName, databaseOpts)
res, err := db.RunCommand(ctx, cmd).DecodeBytes()
// propagate original result
if err == nil {
return bsoncore.Document(res), nil
}
// wrap original error
if mc.bypassSpawn || !strings.Contains(err.Error(), serverSelectionTimeoutStr) {
return nil, MongocryptdError{Wrapped: err}
}
// re-spawn and retry
if err = mc.spawnProcess(); err != nil {
return nil, err
}
res, err = db.RunCommand(ctx, cmd).DecodeBytes()
if err != nil {
return nil, MongocryptdError{Wrapped: err}
}
return bsoncore.Document(res), nil
}
// connect connects the underlying Client instance. This must be called before performing any mark operations.
func (mc *mcryptClient) connect(ctx context.Context) error {
return mc.client.Connect(ctx)
}
// disconnect disconnects the underlying Client instance. This should be called after all operations have completed.
func (mc *mcryptClient) disconnect(ctx context.Context) error {
return mc.client.Disconnect(ctx)
}
func (mc *mcryptClient) spawnProcess() error {
return exec.Command(mc.path, mc.spawnArgs...).Start()
}
// createSpawnArgs creates arguments to spawn mcryptClient. It returns the path and a slice of arguments.
func createSpawnArgs(opts map[string]interface{}) (string, []string) {
var spawnArgs []string
// get command path
path := defaultPath
if p, ok := opts["mongocryptdPath"]; ok {
path = p.(string)
}
// add specified options
if sa, ok := opts["mongocryptdSpawnArgs"]; ok {
spawnArgs = append(spawnArgs, sa.([]string)...)
}
// add timeout options if necessary
var foundTimeout bool
for _, arg := range spawnArgs {
// need to use HasPrefix instead of doing an exact equality check because both
// mongocryptd supports both [--idleShutdownTimeoutSecs, 0] and [--idleShutdownTimeoutSecs=0]
if strings.HasPrefix(arg, "--idleShutdownTimeoutSecs") {
foundTimeout = true
break
}
}
if !foundTimeout {
spawnArgs = append(spawnArgs, defaultTimeoutArgs...)
}
return path, spawnArgs
}