...
1
2
3
4
5
6
7 package mongo
8
9 import (
10 "context"
11 "os/exec"
12 "strings"
13 "time"
14
15 "go.mongodb.org/mongo-driver/mongo/options"
16 "go.mongodb.org/mongo-driver/mongo/readconcern"
17 "go.mongodb.org/mongo-driver/mongo/readpref"
18 "go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
19 )
20
21 const (
22 defaultServerSelectionTimeout = 10 * time.Second
23 defaultURI = "mongodb://localhost:27020"
24 defaultPath = "mongocryptd"
25 serverSelectionTimeoutStr = "server selection error"
26 )
27
28 var defaultTimeoutArgs = []string{"--idleShutdownTimeoutSecs=60"}
29 var databaseOpts = options.Database().SetReadConcern(readconcern.New()).SetReadPreference(readpref.Primary())
30
31 type mongocryptdClient struct {
32 bypassSpawn bool
33 client *Client
34 path string
35 spawnArgs []string
36 }
37
38
39
40
41 func newMongocryptdClient(opts *options.AutoEncryptionOptions) (*mongocryptdClient, error) {
42
43 var bypassSpawn bool
44 var bypassAutoEncryption bool
45
46 if bypass, ok := opts.ExtraOptions["mongocryptdBypassSpawn"]; ok {
47 bypassSpawn = bypass.(bool)
48 }
49 if opts.BypassAutoEncryption != nil {
50 bypassAutoEncryption = *opts.BypassAutoEncryption
51 }
52
53 bypassQueryAnalysis := opts.BypassQueryAnalysis != nil && *opts.BypassQueryAnalysis
54
55 mc := &mongocryptdClient{
56
57
58
59
60 bypassSpawn: bypassSpawn || bypassAutoEncryption || bypassQueryAnalysis,
61 }
62
63 if !mc.bypassSpawn {
64 mc.path, mc.spawnArgs = createSpawnArgs(opts.ExtraOptions)
65 if err := mc.spawnProcess(); err != nil {
66 return nil, err
67 }
68 }
69
70
71 uri := defaultURI
72 if u, ok := opts.ExtraOptions["mongocryptdURI"]; ok {
73 uri = u.(string)
74 }
75
76
77 client, err := NewClient(options.Client().ApplyURI(uri).SetServerSelectionTimeout(defaultServerSelectionTimeout))
78 if err != nil {
79 return nil, err
80 }
81 mc.client = client
82
83 return mc, nil
84 }
85
86
87 func (mc *mongocryptdClient) markCommand(ctx context.Context, dbName string, cmd bsoncore.Document) (bsoncore.Document, error) {
88
89
90
91 ctx = NewSessionContext(ctx, nil)
92 db := mc.client.Database(dbName, databaseOpts)
93
94 res, err := db.RunCommand(ctx, cmd).Raw()
95
96 if err == nil {
97 return bsoncore.Document(res), nil
98 }
99
100 if mc.bypassSpawn || !strings.Contains(err.Error(), serverSelectionTimeoutStr) {
101 return nil, MongocryptdError{Wrapped: err}
102 }
103
104
105 if err = mc.spawnProcess(); err != nil {
106 return nil, err
107 }
108 res, err = db.RunCommand(ctx, cmd).Raw()
109 if err != nil {
110 return nil, MongocryptdError{Wrapped: err}
111 }
112 return bsoncore.Document(res), nil
113 }
114
115
116 func (mc *mongocryptdClient) connect(ctx context.Context) error {
117 return mc.client.Connect(ctx)
118 }
119
120
121 func (mc *mongocryptdClient) disconnect(ctx context.Context) error {
122 return mc.client.Disconnect(ctx)
123 }
124
125 func (mc *mongocryptdClient) spawnProcess() error {
126
127
128 cmd := exec.Command(mc.path, mc.spawnArgs...)
129 cmd.Stdout = nil
130 cmd.Stderr = nil
131 return cmd.Start()
132 }
133
134
135 func createSpawnArgs(opts map[string]interface{}) (string, []string) {
136 var spawnArgs []string
137
138
139 path := defaultPath
140 if p, ok := opts["mongocryptdPath"]; ok {
141 path = p.(string)
142 }
143
144
145 if sa, ok := opts["mongocryptdSpawnArgs"]; ok {
146 spawnArgs = append(spawnArgs, sa.([]string)...)
147 }
148
149
150 var foundTimeout bool
151 for _, arg := range spawnArgs {
152
153
154 if strings.HasPrefix(arg, "--idleShutdownTimeoutSecs") {
155 foundTimeout = true
156 break
157 }
158 }
159 if !foundTimeout {
160 spawnArgs = append(spawnArgs, defaultTimeoutArgs...)
161 }
162
163 return path, spawnArgs
164 }
165
View as plain text