summaryrefslogtreecommitdiff
path: root/src/mongo/rpc/op_msg.h
blob: 88799e568960ddd3616654720b3238588d99ec67 (plain)
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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340

/**
 *    Copyright (C) 2018-present MongoDB, Inc.
 *
 *    This program is free software: you can redistribute it and/or modify
 *    it under the terms of the Server Side Public License, version 1,
 *    as published by MongoDB, Inc.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    Server Side Public License for more details.
 *
 *    You should have received a copy of the Server Side Public License
 *    along with this program. If not, see
 *    <http://www.mongodb.com/licensing/server-side-public-license>.
 *
 *    As a special exception, the copyright holders give permission to link the
 *    code of portions of this program with the OpenSSL library under certain
 *    conditions as described in each individual source file and distribute
 *    linked combinations including the program with the OpenSSL library. You
 *    must comply with the Server Side Public License in all respects for
 *    all of the code used other than as permitted herein. If you modify file(s)
 *    with this exception, you may extend this exception to your version of the
 *    file(s), but you are not obligated to do so. If you do not wish to do so,
 *    delete this exception statement from your version. If you delete this
 *    exception statement from all source files in the program, then also delete
 *    it in the license file.
 */

#pragma once

#include <algorithm>
#include <string>
#include <vector>

#include "mongo/base/string_data.h"
#include "mongo/db/jsobj.h"
#include "mongo/rpc/message.h"

namespace mongo {

struct OpMsg {
    struct DocumentSequence {
        std::string name;
        std::vector<BSONObj> objs;
    };

    static constexpr uint32_t kChecksumPresent = 1 << 0;
    static constexpr uint32_t kMoreToCome = 1 << 1;
    static constexpr uint32_t kExhaustSupported = 1 << 16;

    /**
     * Returns the unvalidated flags for the given message if it is an OP_MSG message.
     * Returns 0 for other message kinds since they are the equivalent of no flags set.
     * Throws if the message is too small to hold flags.
     */
    static uint32_t flags(const Message& message);
    static bool isFlagSet(const Message& message, uint32_t flag) {
        return flags(message) & flag;
    }

    /**
     * Replaces the flags in message with the supplied flags.
     * Only legal on an otherwise valid OP_MSG message.
     */
    static void replaceFlags(Message* message, uint32_t flags);

    /**
     * Adds flag to the list of set flags in message.
     * Only legal on an otherwise valid OP_MSG message.
     */
    static void setFlag(Message* message, uint32_t flag) {
        replaceFlags(message, flags(*message) | flag);
    }

    /**
     * Parses and returns an OpMsg containing unowned BSON.
     */
    static OpMsg parse(const Message& message);

    /**
     * Parses and returns an OpMsg containing owned BSON.
     */
    static OpMsg parseOwned(const Message& message) {
        auto msg = parse(message);
        msg.shareOwnershipWith(message.sharedBuffer());
        return msg;
    }

    Message serialize() const;

    /**
     * Makes all BSONObjs in this object share ownership with buffer.
     */
    void shareOwnershipWith(const ConstSharedBuffer& buffer);

    /**
     * Returns a pointer to the sequence with the given name or nullptr if there are none.
     */
    const DocumentSequence* getSequence(StringData name) const {
        // Getting N sequences is technically O(N**2) but because there currently is at most 2
        // sequences, this does either 1 or 2 comparisons. Consider making sequences a StringMap if
        // there will be many sequences. This problem may also just go away with the IDL project.
        auto it = std::find_if(
            sequences.begin(), sequences.end(), [&](const auto& seq) { return seq.name == name; });
        return it == sequences.end() ? nullptr : &*it;
    }

    BSONObj body;
    std::vector<DocumentSequence> sequences;
};

/**
 * An OpMsg that represents a request. This is a separate type from OpMsg only to provide better
 * type-safety along with a place to hang request-specific methods.
 */
struct OpMsgRequest : public OpMsg {
    // TODO in C++17 remove constructors so we can use aggregate initialization.
    OpMsgRequest() = default;
    explicit OpMsgRequest(OpMsg&& generic) : OpMsg(std::move(generic)) {}

    static OpMsgRequest parse(const Message& message) {
        return OpMsgRequest(OpMsg::parse(message));
    }

    static OpMsgRequest fromDBAndBody(StringData db,
                                      BSONObj body,
                                      const BSONObj& extraFields = {}) {
        OpMsgRequest request;
        request.body = ([&] {
            BSONObjBuilder bodyBuilder(std::move(body));
            bodyBuilder.appendElements(extraFields);
            bodyBuilder.append("$db", db);
            return bodyBuilder.obj();
        }());
        return request;
    }

    StringData getDatabase() const {
        if (auto elem = body["$db"])
            return elem.checkAndGetStringData();
        uasserted(40571, "OP_MSG requests require a $db argument");
    }

    StringData getCommandName() const {
        return body.firstElementFieldName();
    }

    // DO NOT ADD MEMBERS!  Since this type is essentially a strong typedef (see the class comment),
    // it should not hold more data than an OpMsg. It should be freely interconvertible with OpMsg
    // without issues like slicing.
};

/**
 * Builds an OP_MSG message in-place in a Message buffer.
 *
 * While the OP_MSG format imposes no ordering of sections, in order to efficiently support our
 * usage patterns, this class requires that all document sequences (if any) are built before the
 * body. This allows repeatedly appending fields to the body until right before it is ready to be
 * sent.
 */
class OpMsgBuilder {
    MONGO_DISALLOW_COPYING(OpMsgBuilder);

public:
    OpMsgBuilder() {
        skipHeaderAndFlags();
    }

    /**
     * See the documentation for DocSequenceBuilder below.
     */
    class DocSequenceBuilder;
    DocSequenceBuilder beginDocSequence(StringData name);

    /**
     * Returns an empty builder for the body.
     * It is an error to call this if a body has already been begun.  You must destroy or call
     * done() on the returned builder before calling any methods on this object.
     */
    BSONObjBuilder beginBody();
    void setBody(const BSONObj& body) {
        beginBody().appendElements(body);
    }

    /**
     * Returns a builder that can be used to append new fields to the body.
     * It is an error to call this if beginBody() hasn't been called yet. It is an error to append
     * elements with field names that already exist in the body. You must destroy or call done() on
     * the returned builder before calling any methods on this object.
     *
     * TODO decide if it is worth keeping the begin/resume distinction in the public API.
     */
    BSONObjBuilder resumeBody();
    void appendElementsToBody(const BSONObj& body) {
        resumeBody().appendElements(body);
    }

    /**
     * Finish building and return a Message ready to give to the networking layer for transmission.
     * It is illegal to call any methods on this object after calling this.
     */
    Message finish();

    /**
     * Reset this object to its initial empty state. All previously appended data is lost.
     */
    void reset() {
        invariant(!_openBuilder);

        _buf.reset();
        skipHeaderAndFlags();
        _bodyStart = 0;
        _state = kEmpty;
        _openBuilder = false;
    }

    /**
     * Set to true in tests that need to be able to generate duplicate top-level fields to see how
     * the server handles them. Is false by default, although the check only happens in debug
     * builds.
     */
    static AtomicWord<bool> disableDupeFieldCheck_forTest;

    /**
     * Similar to finish, any calls on this object after are illegal.
     */
    BSONObj releaseBody();

    /**
     * Returns whether or not this builder is already building a body.
     */
    bool isBuildingBody() {
        return _state == kBody;
    }

    /**
     * Reserves and claims the bytes requested in the internal BufBuilder.
     */
    void reserveBytes(const std::size_t bytes) {
        _buf.reserveBytes(bytes);
        _buf.claimReservedBytes(bytes);
    }

private:
    friend class DocSequenceBuilder;

    enum State {
        kEmpty,
        kDocSequence,
        kBody,
        kDone,
    };

    void finishDocumentStream(DocSequenceBuilder* docSequenceBuilder);

    void skipHeaderAndFlags() {
        _buf.skip(sizeof(MSGHEADER::Layout));  // This is filled in by finish().
        _buf.appendNum(uint32_t(0));           // flags (currently always 0).
    }

    // When adding members, remember to update reset().
    BufBuilder _buf;
    int _bodyStart = 0;
    State _state = kEmpty;
    bool _openBuilder = false;
};

/**
 * Builds a document sequence in an OpMsgBuilder.
 *
 * Example:
 *
 * auto docSeq = msgBuilder.beginDocSequence("some.sequence");
 *
 * docSeq.append(BSON("a" << 1)); // Copy an obj into the sequence
 *
 * auto bob = docSeq.appendBuilder(); // Build an obj in-place
 * bob.append("a", 2);
 * bob.doneFast();
 *
 * docSeq.done(); // Or just let it go out of scope.
 */
class OpMsgBuilder::DocSequenceBuilder {
    MONGO_DISALLOW_COPYING(DocSequenceBuilder);

public:
    DocSequenceBuilder(DocSequenceBuilder&& other)
        : _buf(other._buf), _msgBuilder(other._msgBuilder), _sizeOffset(other._sizeOffset) {
        other._buf = nullptr;
    }

    ~DocSequenceBuilder() {
        if (_buf)
            done();
    }

    /**
     * Indicates that the caller is done with this stream prior to destruction.
     * Following this call, it is illegal to call any methods on this object.
     */
    void done() {
        invariant(_buf);
        _msgBuilder->finishDocumentStream(this);
        _buf = nullptr;
    }

    /**
     * Appends a single document to this sequence.
     */
    void append(const BSONObj& obj) {
        _buf->appendBuf(obj.objdata(), obj.objsize());
    }

    /**
     * Returns a BSONObjBuilder that appends a single document to this sequence in place.
     * It is illegal to call any methods on this DocSequenceBuilder until the returned builder
     * is destroyed or done()/doneFast() is called on it.
     */
    BSONObjBuilder appendBuilder() {
        return BSONObjBuilder(*_buf);
    }

    int len() const {
        return _buf->len();
    }

private:
    friend OpMsgBuilder;

    DocSequenceBuilder(OpMsgBuilder* msgBuilder, BufBuilder* buf, int sizeOffset)
        : _buf(buf), _msgBuilder(msgBuilder), _sizeOffset(sizeOffset) {}

    BufBuilder* _buf;
    OpMsgBuilder* const _msgBuilder;
    const int _sizeOffset;
};

}  // namespace mongo