xxq250
/
homework-jianmu
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

[TD-2101]<fix> prevent buf overflow when write insert sql

This commit is contained in:
Hui Li 2020-11-16 19:55:25 +08:00
parent f9eeaa1b19
commit 021ea48b55
1 changed files with 46 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

74
src/kit/taosdump/taosdump.c
View File

@ -182,6 +182,7 @@ static struct argp_option options[] = {
{"start-time", 'S', "START_TIME", 0, "Start time to dump.", 3}, {"start-time", 'S', "START_TIME", 0, "Start time to dump.", 3},
{"end-time", 'E', "END_TIME", 0, "End time to dump.", 3}, {"end-time", 'E', "END_TIME", 0, "End time to dump.", 3},
{"data-batch", 'N', "DATA_BATCH", 0, "Number of data point per insert statement. Default is 1.", 3}, {"data-batch", 'N', "DATA_BATCH", 0, "Number of data point per insert statement. Default is 1.", 3},
{"max-sql-len", 'L', "SQL_LEN", 0, "Max length of one sql. Default is 65480.", 3},
{"table-batch", 't', "TABLE_BATCH", 0, "Number of table dumpout into one output file. Default is 1.", 3}, {"table-batch", 't', "TABLE_BATCH", 0, "Number of table dumpout into one output file. Default is 1.", 3},
{"thread_num", 'T', "THREAD_NUM", 0, "Number of thread for dump in file. Default is 5.", 3}, {"thread_num", 'T', "THREAD_NUM", 0, "Number of thread for dump in file. Default is 5.", 3},
{"allow-sys", 'a', 0, 0, "Allow to dump sys database", 3}, {"allow-sys", 'a', 0, 0, "Allow to dump sys database", 3},
@ -209,6 +210,7 @@ struct arguments {
int64_t start_time; int64_t start_time;
int64_t end_time; int64_t end_time;
int32_t data_batch; int32_t data_batch;
int32_t max_sql_len;
int32_t table_batch; // num of table which will be dump into one output file. int32_t table_batch; // num of table which will be dump into one output file.
bool allow_sys; bool allow_sys;
// other options // other options
@ -307,6 +309,17 @@ static error_t parse_opt(int key, char *arg, struct argp_state *state) {
case 'N': case 'N':
arguments->data_batch = atoi(arg); arguments->data_batch = atoi(arg);
break; break;
case 'L':
{
int32_t len = atoi(arg);
if (len > TSDB_MAX_ALLOWED_SQL_LEN) {
len = TSDB_MAX_ALLOWED_SQL_LEN;
} else if (len < TSDB_MAX_SQL_LEN) {
len = TSDB_MAX_SQL_LEN;
}
arguments->max_sql_len = len;
break;
}
case 't': case 't':
arguments->table_batch = atoi(arg); arguments->table_batch = atoi(arg);
break; break;
@ -369,6 +382,7 @@ struct arguments tsArguments = {
0, 0,
INT64_MAX, INT64_MAX,
1, 1,
TSDB_MAX_SQL_LEN,
1, 1,
false, false,
// other options // other options
@ -424,6 +438,7 @@ int main(int argc, char *argv[]) {
printf("start_time: %" PRId64 "\n", tsArguments.start_time); printf("start_time: %" PRId64 "\n", tsArguments.start_time);
printf("end_time: %" PRId64 "\n", tsArguments.end_time); printf("end_time: %" PRId64 "\n", tsArguments.end_time);
printf("data_batch: %d\n", tsArguments.data_batch); printf("data_batch: %d\n", tsArguments.data_batch);
printf("max_sql_len: %d\n", tsArguments.max_sql_len);
printf("table_batch: %d\n", tsArguments.table_batch); printf("table_batch: %d\n", tsArguments.table_batch);
printf("thread_num: %d\n", tsArguments.thread_num); printf("thread_num: %d\n", tsArguments.thread_num);
printf("allow_sys: %d\n", tsArguments.allow_sys); printf("allow_sys: %d\n", tsArguments.allow_sys);
@ -1479,7 +1494,8 @@ int taosDumpTableData(FILE *fp, char *tbname, struct arguments *arguments, TAOS*
return -1; return -1;
} }
char* tmpBuffer = (char *)calloc(1, COMMAND_SIZE); int32_t sql_buf_len = arguments->max_sql_len;
char* tmpBuffer = (char *)calloc(1, sql_buf_len + 128);
if (tmpBuffer == NULL) { if (tmpBuffer == NULL) {
fprintf(stderr, "failed to allocate memory\n"); fprintf(stderr, "failed to allocate memory\n");
free(tmpCommand); free(tmpCommand);
@ -1522,85 +1538,83 @@ int taosDumpTableData(FILE *fp, char *tbname, struct arguments *arguments, TAOS*
return -1; return -1;
} }
char sqlStr[8] = "\0";
if (arguments->mysqlFlag) {
sprintf(sqlStr, "INSERT");
} else {
sprintf(sqlStr, "IMPORT");
}
int rowFlag = 0; int rowFlag = 0;
int32_t curr_sqlstr_len = 0;
int32_t total_sqlstr_len = 0;
count = 0; count = 0;
while ((row = taos_fetch_row(tmpResult)) != NULL) { while ((row = taos_fetch_row(tmpResult)) != NULL) {
pstr = tmpBuffer; pstr = tmpBuffer;
curr_sqlstr_len = 0;
int32_t* length = taos_fetch_lengths(tmpResult); // act len int32_t* length = taos_fetch_lengths(tmpResult); // act len
if (count == 0) { if (count == 0) {
pstr += sprintf(pstr, "%s INTO %s VALUES (", sqlStr, tbname); total_sqlstr_len = 0;
curr_sqlstr_len += sprintf(pstr + curr_sqlstr_len, "INSERT INTO %s VALUES (", tbname);
} else { } else {
if (arguments->mysqlFlag) { if (arguments->mysqlFlag) {
if (0 == rowFlag) { if (0 == rowFlag) {
pstr += sprintf(pstr, "("); curr_sqlstr_len += sprintf(pstr + curr_sqlstr_len, "(");
rowFlag++; rowFlag++;
} else { } else {
pstr += sprintf(pstr, ", ("); curr_sqlstr_len += sprintf(pstr + curr_sqlstr_len, ", (");
} }
} else { } else {
pstr += sprintf(pstr, "("); curr_sqlstr_len += sprintf(pstr + curr_sqlstr_len, "(");
} }
} }
for (int col = 0; col < numFields; col++) { for (int col = 0; col < numFields; col++) {
if (col != 0) pstr += sprintf(pstr, ", "); if (col != 0) curr_sqlstr_len += sprintf(pstr + curr_sqlstr_len, ", ");
if (row[col] == NULL) { if (row[col] == NULL) {
pstr += sprintf(pstr, "NULL"); curr_sqlstr_len += sprintf(pstr + curr_sqlstr_len, "NULL");
continue; continue;
} }
switch (fields[col].type) { switch (fields[col].type) {
case TSDB_DATA_TYPE_BOOL: case TSDB_DATA_TYPE_BOOL:
pstr += sprintf(pstr, "%d", ((((int32_t)(*((char *)row[col]))) == 1) ? 1 : 0)); curr_sqlstr_len += sprintf(pstr + curr_sqlstr_len, "%d", ((((int32_t)(*((char *)row[col]))) == 1) ? 1 : 0));
break; break;
case TSDB_DATA_TYPE_TINYINT: case TSDB_DATA_TYPE_TINYINT:
pstr += sprintf(pstr, "%d", *((int8_t *)row[col])); curr_sqlstr_len += sprintf(pstr + curr_sqlstr_len, "%d", *((int8_t *)row[col]));
break; break;
case TSDB_DATA_TYPE_SMALLINT: case TSDB_DATA_TYPE_SMALLINT:
pstr += sprintf(pstr, "%d", *((int16_t *)row[col])); curr_sqlstr_len += sprintf(pstr + curr_sqlstr_len, "%d", *((int16_t *)row[col]));
break; break;
case TSDB_DATA_TYPE_INT: case TSDB_DATA_TYPE_INT:
pstr += sprintf(pstr, "%d", *((int32_t *)row[col])); curr_sqlstr_len += sprintf(pstr + curr_sqlstr_len, "%d", *((int32_t *)row[col]));
break; break;
case TSDB_DATA_TYPE_BIGINT: case TSDB_DATA_TYPE_BIGINT:
pstr += sprintf(pstr, "%" PRId64 "", *((int64_t *)row[col])); curr_sqlstr_len += sprintf(pstr + curr_sqlstr_len, "%" PRId64 "", *((int64_t *)row[col]));
break; break;
case TSDB_DATA_TYPE_FLOAT: case TSDB_DATA_TYPE_FLOAT:
pstr += sprintf(pstr, "%f", GET_FLOAT_VAL(row[col])); curr_sqlstr_len += sprintf(pstr + curr_sqlstr_len, "%f", GET_FLOAT_VAL(row[col]));
break; break;
case TSDB_DATA_TYPE_DOUBLE: case TSDB_DATA_TYPE_DOUBLE:
pstr += sprintf(pstr, "%f", GET_DOUBLE_VAL(row[col])); curr_sqlstr_len += sprintf(pstr + curr_sqlstr_len, "%f", GET_DOUBLE_VAL(row[col]));
break; break;
case TSDB_DATA_TYPE_BINARY: case TSDB_DATA_TYPE_BINARY:
*(pstr++) = '\''; //*(pstr++) = '\'';
converStringToReadable((char *)row[col], length[col], tbuf, COMMAND_SIZE); converStringToReadable((char *)row[col], length[col], tbuf, COMMAND_SIZE);
pstr = stpcpy(pstr, tbuf); //pstr = stpcpy(pstr, tbuf);
*(pstr++) = '\''; //*(pstr++) = '\'';
pstr += sprintf(pstr + curr_sqlstr_len, "\'%s\'", tbuf);
break; break;
case TSDB_DATA_TYPE_NCHAR: case TSDB_DATA_TYPE_NCHAR:
convertNCharToReadable((char *)row[col], length[col], tbuf, COMMAND_SIZE); convertNCharToReadable((char *)row[col], length[col], tbuf, COMMAND_SIZE);
pstr += sprintf(pstr, "\'%s\'", tbuf); pstr += sprintf(pstr + curr_sqlstr_len, "\'%s\'", tbuf);
break; break;
case TSDB_DATA_TYPE_TIMESTAMP: case TSDB_DATA_TYPE_TIMESTAMP:
if (!arguments->mysqlFlag) { if (!arguments->mysqlFlag) {
pstr += sprintf(pstr, "%" PRId64 "", *(int64_t *)row[col]); curr_sqlstr_len += sprintf(pstr + curr_sqlstr_len, "%" PRId64 "", *(int64_t *)row[col]);
} else { } else {
char buf[64] = "\0"; char buf[64] = "\0";
int64_t ts = *((int64_t *)row[col]); int64_t ts = *((int64_t *)row[col]);
time_t tt = (time_t)(ts / 1000); time_t tt = (time_t)(ts / 1000);
struct tm *ptm = localtime(&tt); struct tm *ptm = localtime(&tt);
strftime(buf, 64, "%y-%m-%d %H:%M:%S", ptm); strftime(buf, 64, "%y-%m-%d %H:%M:%S", ptm);
pstr += sprintf(pstr, "\'%s.%03d\'", buf, (int)(ts % 1000)); curr_sqlstr_len += sprintf(pstr + curr_sqlstr_len, "\'%s.%03d\'", buf, (int)(ts % 1000));
} }
break; break;
default: default:
@ -1608,13 +1622,15 @@ int taosDumpTableData(FILE *fp, char *tbname, struct arguments *arguments, TAOS*
} }
} }
pstr += sprintf(pstr, ") "); curr_sqlstr_len += sprintf(pstr + curr_sqlstr_len, ") ");
totalRows++; totalRows++;
count++; count++;
fprintf(fp, "%s", tmpBuffer); fprintf(fp, "%s", tmpBuffer);
if (count >= arguments->data_batch) { total_sqlstr_len += curr_sqlstr_len;
if ((count >= arguments->data_batch) || (sql_buf_len - total_sqlstr_len < TSDB_MAX_BYTES_PER_ROW)) {
fprintf(fp, ";\n"); fprintf(fp, ";\n");
count = 0; count = 0;
} //else { } //else {