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
|
/*
* Copyright 2009-2015 MongoDB, Inc.
*
* 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
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
/* Include Python.h so we can set Python's error indicator. */
#define PY_SSIZE_T_CLEAN
#include "Python.h"
#include <stdlib.h>
#include <string.h>
#include "buffer.h"
#define INITIAL_BUFFER_SIZE 256
struct buffer {
char* buffer;
int size;
int position;
};
/* Set Python's error indicator to MemoryError.
* Called after allocation failures. */
static void set_memory_error(void) {
PyErr_NoMemory();
}
/* Allocate and return a new buffer.
* Return NULL and sets MemoryError on allocation failure. */
buffer_t pymongo_buffer_new(void) {
buffer_t buffer;
buffer = (buffer_t)malloc(sizeof(struct buffer));
if (buffer == NULL) {
set_memory_error();
return NULL;
}
buffer->size = INITIAL_BUFFER_SIZE;
buffer->position = 0;
buffer->buffer = (char*)malloc(sizeof(char) * INITIAL_BUFFER_SIZE);
if (buffer->buffer == NULL) {
free(buffer);
set_memory_error();
return NULL;
}
return buffer;
}
/* Free the memory allocated for `buffer`.
* Return non-zero on failure. */
int pymongo_buffer_free(buffer_t buffer) {
if (buffer == NULL) {
return 1;
}
/* Buffer will be NULL when buffer_grow fails. */
if (buffer->buffer != NULL) {
free(buffer->buffer);
}
free(buffer);
return 0;
}
/* Grow `buffer` to at least `min_length`.
* Return non-zero and sets MemoryError on allocation failure. */
static int buffer_grow(buffer_t buffer, int min_length) {
int old_size = 0;
int size = buffer->size;
char* old_buffer = buffer->buffer;
if (size >= min_length) {
return 0;
}
while (size < min_length) {
old_size = size;
size *= 2;
if (size <= old_size) {
/* Size did not increase. Could be an overflow
* or size < 1. Just go with min_length. */
size = min_length;
}
}
buffer->buffer = (char*)realloc(buffer->buffer, sizeof(char) * size);
if (buffer->buffer == NULL) {
free(old_buffer);
set_memory_error();
return 1;
}
buffer->size = size;
return 0;
}
/* Assure that `buffer` has at least `size` free bytes (and grow if needed).
* Return non-zero and sets MemoryError on allocation failure.
* Return non-zero and sets ValueError if `size` would exceed 2GiB. */
static int buffer_assure_space(buffer_t buffer, int size) {
int new_size = buffer->position + size;
/* Check for overflow. */
if (new_size < buffer->position) {
PyErr_SetString(PyExc_ValueError,
"Document would overflow BSON size limit");
return 1;
}
if (new_size <= buffer->size) {
return 0;
}
return buffer_grow(buffer, new_size);
}
/* Save `size` bytes from the current position in `buffer` (and grow if needed).
* Return offset for writing, or -1 on failure.
* Sets MemoryError or ValueError on failure. */
buffer_position pymongo_buffer_save_space(buffer_t buffer, int size) {
int position = buffer->position;
if (buffer_assure_space(buffer, size) != 0) {
return -1;
}
buffer->position += size;
return position;
}
/* Write `size` bytes from `data` to `buffer` (and grow if needed).
* Return non-zero on failure.
* Sets MemoryError or ValueError on failure. */
int pymongo_buffer_write(buffer_t buffer, const char* data, int size) {
if (buffer_assure_space(buffer, size) != 0) {
return 1;
}
memcpy(buffer->buffer + buffer->position, data, size);
buffer->position += size;
return 0;
}
int pymongo_buffer_get_position(buffer_t buffer) {
return buffer->position;
}
char* pymongo_buffer_get_buffer(buffer_t buffer) {
return buffer->buffer;
}
void pymongo_buffer_update_position(buffer_t buffer, buffer_position new_position) {
buffer->position = new_position;
}
|
