From: teodor <teodor>
Date: Thu, 29 Nov 2007 18:17:32 +0000 (+0000)
Subject: Add max allowed size to TC_Read and TC_Talk
X-Git-Url: http://sigaev.ru/git/gitweb.cgi?a=commitdiff_plain;h=2405d79bcae81049b002b8913060b840370f9a11;p=tedtools.git

Add max allowed size to TC_Read and TC_Talk
---

diff --git a/connection.h b/connection.h
index 7c0db10..6ccd725 100644
--- a/connection.h
+++ b/connection.h
@@ -81,8 +81,8 @@ u_int32_t TC_ClientInitConnection(TC_Connection *cs, char *name, u_int32_t port)
 u_int32_t TC_ServerInitConnect( TC_Connection *cs );
 u_int32_t TC_ServerConnect( TC_Connection *cs, int timeout );
 u_int32_t TC_Send( TC_Connection *cs );
-u_int32_t TC_Read( TC_Connection *cs );
-u_int32_t TC_Talk( TC_Connection *cs );
+u_int32_t TC_Read( TC_Connection *cs, size_t maxsize );
+u_int32_t TC_Talk( TC_Connection *cs, size_t maxsize );
 void TC_FreeConnection( TC_Connection *cs );
 int TC_ReadyIO( TC_Connection **cs, int number, int timeout );
 
diff --git a/tcp.c b/tcp.c
index fe59586..7183b42 100644
--- a/tcp.c
+++ b/tcp.c
@@ -394,7 +394,7 @@ resizeCS( TC_Connection *cs, int sz ) {
 }
 
 u_int32_t
-TC_Read( TC_Connection *cs ) {
+TC_Read( TC_Connection *cs, size_t maxsize ) {
 	int sz, totalread = -1, toread=0, alreadyread;
 
 	if ( cs->state == CS_ERROR )
@@ -411,6 +411,12 @@ TC_Read( TC_Connection *cs ) {
 		resizeCS(cs, sizeof(u_int32_t));
 	} else {
 		totalread = *(u_int32_t*)(cs->buf);
+		if ( maxsize > 0 && totalread > maxsize )
+		{
+			tlog(TL_ALARM,"TC_Read: message size (%d b) is greater than max allowed (%d b)", totalread, maxsize);
+			cs->state = CS_ERROR;
+			return CS_ERROR;
+		}
 		toread = totalread - alreadyread;
 		if ( toread == 0 ) {
 			cs->state = CS_FINISHREAD;
@@ -458,7 +464,7 @@ TC_FreeConnection( TC_Connection *cs ) {
 }
 
 u_int32_t 
-TC_Talk( TC_Connection *cs ) {
+TC_Talk( TC_Connection *cs, size_t maxsize  ) {
 	if ( cs->state==CS_NOTINITED ) 
 		TC_ServerInitConnect( cs );
 
@@ -479,7 +485,7 @@ TC_Talk( TC_Connection *cs ) {
 	cs->ptr = cs->buf;
 	while( cs->state != CS_FINISHREAD ) {
 		while( !TC_ReadyIO( &cs, 1, 100) );
-		if ( TC_Read(cs) == CS_ERROR ) return CS_ERROR;
+		if ( TC_Read(cs, maxsize) == CS_ERROR ) return CS_ERROR;
 	}
 
 	return CS_OK;