diff --git a/client.c b/client.c index 2398efb..6f2c29e 100644 --- a/client.c +++ b/client.c @@ -31,6 +31,8 @@ static uint32_t file_size; static int recv_bin = 0; static config_t conf; +#define DEBUG + void getMessage(void *p) { int c = 0, offset = 0; cmdpacket_t cp = (cmdpacket_t)bufr; @@ -55,31 +57,36 @@ void getMessage(void *p) { #ifdef DEBUG printf("raw data len: %d\n", datalen); #endif - char* data = malloc(datalen); - offset = c - ++i; - if(offset > 0) { - memcpy(data, bufr+i, offset); - #ifdef DEBUG - printf("copy %d bytes data that had been received.\n", offset); - #endif + if(datalen == 0) { + puts("raw data is empty, truncate file."); + fclose(fopen(savepath, "wb+")); + } else { + char* data = malloc(datalen); + offset = c - ++i; + if(offset > 0) { + memcpy(data, bufr+i, offset); + #ifdef DEBUG + printf("copy %d bytes data that had been received.\n", offset); + #endif + } + else offset = 0; + if(datalen-offset == recv(sockfd, data+offset, datalen-offset, MSG_WAITALL)) { + //FILE* fp = fopen("raw_before_dec", "wb+"); + //fwrite(data, datalen, 1, fp); + //fclose(fp); + off_t tmp = datalen; + char* newdata = raw_decrypt(data, &tmp, 0, conf.pwd); + if(newdata) { + printf("raw data len after decode: %d\n", (int)tmp); + FILE* fp = fopen(savepath, "wb+"); + fwrite(newdata, (size_t)tmp, 1, fp); + fclose(fp); + free(newdata); + puts("recv raw data succeed."); + } else puts("decode raw data error."); + } else puts("recv raw data error."); + free(data); } - else offset = 0; - if(datalen-offset == recv(sockfd, data+offset, datalen-offset, MSG_WAITALL)) { - //FILE* fp = fopen("raw_before_dec", "wb+"); - //fwrite(data, datalen, 1, fp); - //fclose(fp); - off_t tmp = datalen; - char* newdata = raw_decrypt(data, &tmp, 0, conf.pwd); - if(newdata) { - printf("raw data len after decode: %d\n", (int)tmp); - FILE* fp = fopen(savepath, "wb+"); - fwrite(newdata, (size_t)tmp, 1, fp); - fclose(fp); - free(newdata); - puts("recv raw data succeed."); - } else puts("decode raw data error."); - } else puts("recv raw data error."); - free(data); recv_bin = offset = 0; } } else { @@ -210,37 +217,32 @@ int main(int argc, char *argv[]) { // usage: ./client cfg.sp host port if(!strcmp(buf, "set")) { p->cmd = CMDSET; p->datalen = strlen(buf+4); - memcpy(p->data, buf+4, p->datalen); - cmdpacket_encrypt(p, 0, conf.sps); + cmdpacket_encrypt(p, 0, conf.sps, buf+4); send_cmd(sockfd, p); free(p); } else if(!strcmp(buf, "dat")) { p->cmd = CMDDAT; p->datalen = strlen(buf+4); - memcpy(p->data, buf+4, p->datalen); - cmdpacket_encrypt(p, 0, conf.sps); + cmdpacket_encrypt(p, 0, conf.sps, buf+4); send_cmd(sockfd, p); free(p); } else if(!strcmp(buf, "get")) { p->cmd = CMDGET; p->datalen = strlen(buf+4); - memcpy(p->data, buf+4, p->datalen); - cmdpacket_encrypt(p, 0, conf.pwd); + cmdpacket_encrypt(p, 0, conf.pwd, (const char *)&buf+4); send_cmd(sockfd, p); free(p); } else if(!strcmp(buf, "cat")) { p->cmd = CMDCAT; p->datalen = 4; - memcpy(p->data, "fill", p->datalen); recv_bin = 1; - cmdpacket_encrypt(p, 0, conf.pwd); + cmdpacket_encrypt(p, 0, conf.pwd, (const char *)&"fill"); send_cmd(sockfd, p); free(p); } else if(!strcmp(buf, "del")) { p->cmd = CMDDEL; p->datalen = strlen(buf+4); - memcpy(p->data, buf+4, p->datalen); - cmdpacket_encrypt(p, 0, conf.sps); + cmdpacket_encrypt(p, 0, conf.sps, buf+4); send_cmd(sockfd, p); free(p); } else if(!strcmp(buf, "md5")) { @@ -252,15 +254,14 @@ int main(int argc, char *argv[]) { // usage: ./client cfg.sp host port printf("Read md5:"); for(int i = 0; i < 16; i++) printf("%02x", (uint8_t)(p->data[i])); putchar('\n'); - cmdpacket_encrypt(p, 0, conf.pwd); + cmdpacket_encrypt(p, 0, conf.pwd, (const char *)&p->data); send_cmd(sockfd, p); free(p); } } else if(!strcmp(buf, "end")) { p->cmd = CMDEND; p->datalen = 4; - memcpy(p->data, "fill", p->datalen); - cmdpacket_encrypt(p, 0, conf.pwd); + cmdpacket_encrypt(p, 0, conf.pwd, (const char *)&"fill"); send_cmd(sockfd, p); free(p); exit(EXIT_SUCCESS); diff --git a/crypto.h b/crypto.h index 4a1c6d4..1ff697b 100644 --- a/crypto.h +++ b/crypto.h @@ -93,13 +93,13 @@ static char* raw_decrypt(const char* buf, off_t* len, int index, const char pwd[ return decbuf; } -static void cmdpacket_encrypt(cmdpacket_t p, int index, const char pwd[64]) { - TEADAT tin = {p->datalen, p->data}; +static void cmdpacket_encrypt(cmdpacket_t p, int index, const char pwd[64], const char* data) { + TEADAT tin = {p->datalen, (uint8_t *)data}; TEADAT tout; TEA tea[4]; #ifdef DEBUG printf("encrypt len: %d, data: ", p->datalen); - for(int i = 0; i < p->datalen; i++) printf("%02x", p->data[i]); + for(int i = 0; i < p->datalen; i++) printf("%02x", data[i]); putchar('\n'); #endif @@ -115,7 +115,7 @@ static void cmdpacket_encrypt(cmdpacket_t p, int index, const char pwd[64]) { tea_encrypt_native_endian(tea, sumtable, &tin, &tout); - md5(p->data, p->datalen, p->md5); + md5((const uint8_t *)data, p->datalen, p->md5); #ifdef DEBUG printf("encrypt md5: "); for(int i = 0; i < 16; i++) printf("%02x", p->md5[i]); diff --git a/server.c b/server.c index 3e5f3f9..e47c273 100644 --- a/server.c +++ b/server.c @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -125,8 +126,7 @@ static int send_data(int accept_fd, int index, server_ack_t cmd, const char *dat cmdpacket_t p = (cmdpacket_t)buf; p->cmd = (uint8_t)cmd; p->datalen = length; - memcpy(p->data, data, p->datalen); - cmdpacket_encrypt(p, index, cfg.pwd); + cmdpacket_encrypt(p, index, cfg.pwd, data); int total = CMDPACKET_HEAD_LEN+p->datalen; if(!~send(accept_fd, buf, total, 0)) { perror("Send data error"); @@ -145,8 +145,18 @@ static int send_all(thread_timer_t *timer) { if(fp == NULL) return 1; pthread_cleanup_push((void*)&close_dict, (void*)(uintptr_t)timer->index); off_t len = 0, file_size = get_dict_size(); - char* buf = (char*)malloc(file_size); - if(buf) { + + while(1) { + if(file_size <= 0) { + re = send(timer->accept_fd, "0$0123456789abcdef", CMDPACKET_HEAD_LEN, 0); + puts("Send 0 bytes."); + break; + } + char* buf = (char*)malloc(file_size); + if(!buf) { + perror("malloc"); + break; + } pthread_cleanup_push((void*)&free, (void*)buf); if(fread(buf, file_size, 1, fp) == 1) { #ifdef DEBUG @@ -154,19 +164,16 @@ static int send_all(thread_timer_t *timer) { #endif char* encbuf = raw_encrypt(buf, &file_size, timer->index, cfg.pwd); sprintf(timer->dat, "%u$", (unsigned int)file_size); - //printf("Get encrypted file size: %s\n", timer->dat); - //FILE* fp = fopen("raw_after_enc", "wb+"); - //fwrite(encbuf, file_size, 1, fp); - //fclose(fp); pthread_cleanup_push((void*)&free, (void*)encbuf); - if(send(timer->accept_fd, timer->dat, strlen(timer->dat), 0) > 0) { - re = send(timer->accept_fd, encbuf, file_size, 0); - printf("Send %u bytes.\n", re); - } else re = 0; + struct iovec iov[2] = {{timer->dat, strlen(timer->dat)}, {encbuf, file_size}}; + re = writev(timer->accept_fd, (const struct iovec *)&iov, 2); + printf("Send %d bytes.\n", re); pthread_cleanup_pop(1); } pthread_cleanup_pop(1); + break; } + pthread_cleanup_pop(1); return re; }