diff --git a/doc/nn_getsockopt.txt b/doc/nn_getsockopt.txt index d0b10aa98..c47521aab 100644 --- a/doc/nn_getsockopt.txt +++ b/doc/nn_getsockopt.txt @@ -47,6 +47,10 @@ __ header defines generic socket-level options larger than the buffer, exactly one message may be buffered in addition to the data in the receive buffer. The type of this option is int. Default value is 128kB. +*NN_RCVMAXSIZE*:: + Maximum message size that can be received, in bytes. Negative value means + that the received size is limited only by available addressable memory. The + type of this option is int. Default is 1024kB. *NN_SNDTIMEO*:: The timeout for send operation on the socket, in milliseconds. If message cannot be sent within the specified timeout, EAGAIN error is returned. diff --git a/doc/nn_setsockopt.txt b/doc/nn_setsockopt.txt index 3abe93675..be037f7f6 100644 --- a/doc/nn_setsockopt.txt +++ b/doc/nn_setsockopt.txt @@ -41,6 +41,10 @@ __ header defines generic socket-level options larger than the buffer, exactly one message may be buffered in addition to the data in the receive buffer. The type of this option is int. Default value is 128kB. +*NN_RCVMAXSIZE*:: + Maximum message size that can be received, in bytes. Negative value means + that the received size is limited only by available addressable memory. The + type of this option is int. Default is 1024kB. *NN_SNDTIMEO*:: The timeout for send operation on the socket, in milliseconds. If message cannot be sent within the specified timeout, EAGAIN error is returned. diff --git a/src/core/sock.c b/src/core/sock.c index 546d2f349..f01b9b158 100644 --- a/src/core/sock.c +++ b/src/core/sock.c @@ -122,6 +122,7 @@ int nn_sock_init (struct nn_sock *self, struct nn_socktype *socktype, int fd) self->linger = 1000; self->sndbuf = 128 * 1024; self->rcvbuf = 128 * 1024; + self->rcvmaxsize = 1024 * 1024; self->sndtimeo = -1; self->rcvtimeo = -1; self->reconnect_ivl = 100; @@ -313,6 +314,11 @@ static int nn_sock_setopt_inner (struct nn_sock *self, int level, return -EINVAL; dst = &self->rcvbuf; break; + case NN_RCVMAXSIZE: + if (nn_slow (val < -1)) + return -EINVAL; + dst = &self->rcvmaxsize; + break; case NN_SNDTIMEO: dst = &self->sndtimeo; break; @@ -397,6 +403,9 @@ int nn_sock_getopt_inner (struct nn_sock *self, int level, case NN_RCVBUF: intval = self->rcvbuf; break; + case NN_RCVMAXSIZE: + intval = self->rcvmaxsize; + break; case NN_SNDTIMEO: intval = self->sndtimeo; break; diff --git a/src/core/sock.h b/src/core/sock.h index 2e6687f0e..82e57ef74 100644 --- a/src/core/sock.h +++ b/src/core/sock.h @@ -82,6 +82,7 @@ struct nn_sock int linger; int sndbuf; int rcvbuf; + int rcvmaxsize; int sndtimeo; int rcvtimeo; int reconnect_ivl; diff --git a/src/core/symbol.c b/src/core/symbol.c index 22ba26211..efe474570 100644 --- a/src/core/symbol.c +++ b/src/core/symbol.c @@ -136,6 +136,8 @@ static const struct nn_symbol_properties sym_value_names [] = { NN_TYPE_INT, NN_UNIT_BYTES}, {NN_RCVBUF, "NN_RCVBUF", NN_NS_SOCKET_OPTION, NN_TYPE_INT, NN_UNIT_BYTES}, + {NN_RCVMAXSIZE, "NN_RCVMAXSIZE", NN_NS_SOCKET_OPTION, + NN_TYPE_INT, NN_UNIT_BYTES}, {NN_SNDTIMEO, "NN_SNDTIMEO", NN_NS_SOCKET_OPTION, NN_TYPE_INT, NN_UNIT_MILLISECONDS}, {NN_RCVTIMEO, "NN_RCVTIMEO", NN_NS_SOCKET_OPTION, diff --git a/src/nn.h b/src/nn.h index 4fa42c5e8..3d508d932 100644 --- a/src/nn.h +++ b/src/nn.h @@ -336,6 +336,7 @@ NN_EXPORT struct nn_cmsghdr *nn_cmsg_nxthdr_ ( #define NN_PROTOCOL 13 #define NN_IPV4ONLY 14 #define NN_SOCKET_NAME 15 +#define NN_RCVMAXSIZE 16 /* Send/recv options. */ #define NN_DONTWAIT 1 diff --git a/src/transports/tcp/stcp.c b/src/transports/tcp/stcp.c index 3625898f8..a19236a3c 100644 --- a/src/transports/tcp/stcp.c +++ b/src/transports/tcp/stcp.c @@ -204,6 +204,8 @@ static void nn_stcp_handler (struct nn_fsm *self, int src, int type, int rc; struct nn_stcp *stcp; uint64_t size; + int opt; + size_t opt_sz = sizeof (opt); stcp = nn_cont (self, struct nn_stcp, fsm); @@ -322,9 +324,21 @@ static void nn_stcp_handler (struct nn_fsm *self, int src, int type, switch (stcp->instate) { case NN_STCP_INSTATE_HDR: - /* Message header was received. Allocate memory for the - message. */ + /* Message header was received. Check that message size + is acceptable by comparing with NN_RCVMAXSIZE; + if it's too large, drop the connection. */ size = nn_getll (stcp->inhdr); + + nn_pipebase_getopt (&stcp->pipebase, NN_SOL_SOCKET, + NN_RCVMAXSIZE, &opt, &opt_sz); + + if (opt != -1 && size > opt) { + stcp->state = NN_STCP_STATE_DONE; + nn_fsm_raise (&stcp->fsm, &stcp->done, NN_STCP_ERROR); + return; + } + + /* Allocate memory for the message. */ nn_msg_term (&stcp->inmsg); nn_msg_init (&stcp->inmsg, (size_t) size); diff --git a/tests/tcp.c b/tests/tcp.c index 50a4d663b..a91bb9b93 100644 --- a/tests/tcp.c +++ b/tests/tcp.c @@ -41,6 +41,7 @@ int main () int opt; size_t sz; int s1, s2; + void * dummy_buf; /* Try closing bound but unconnected socket. */ sb = test_socket (AF_SP, NN_PAIR); @@ -178,6 +179,35 @@ int main () test_close (sc); test_close (s1); + /* Test NN_RCVMAXSIZE limit */ + sb = test_socket (AF_SP, NN_PAIR); + test_bind (sb, SOCKET_ADDRESS); + s1 = test_socket (AF_SP, NN_PAIR); + test_connect (s1, SOCKET_ADDRESS); + opt = 4; + rc = nn_setsockopt (sb, NN_SOL_SOCKET, NN_RCVMAXSIZE, &opt, sizeof (opt)); + nn_assert (rc == 0); + nn_sleep (100); + test_send (s1, "ABC"); + test_recv (sb, "ABC"); + test_send (s1, "0123456789012345678901234567890123456789"); + rc = nn_recv (sb, dummy_buf, NN_MSG, NN_DONTWAIT); + nn_assert (rc < 0); + errno_assert (nn_errno () == EAGAIN); + test_close (sb); + test_close (s1); + + /* Test that NN_RCVMAXSIZE can be -1, but not lower */ + sb = test_socket (AF_SP, NN_PAIR); + opt = -1; + rc = nn_setsockopt (sb, NN_SOL_SOCKET, NN_RCVMAXSIZE, &opt, sizeof (opt)); + nn_assert (rc >= 0); + opt = -2; + rc = nn_setsockopt (sb, NN_SOL_SOCKET, NN_RCVMAXSIZE, &opt, sizeof (opt)); + nn_assert (rc < 0); + errno_assert (nn_errno () == EINVAL); + test_close (sb); + return 0; }