Add max allowed size to TC_Read and TC_Talk

diff --git a/connection.h b/connection.h
index 7c0db10..6ccd725 100644 (file)
--- 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 (file)
--- 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;