Add max allowed size to TC_Read and TC_Talk

[tedtools.git] / tcp.c

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;