nbd-server: Limit the number of connections

[nbd.git] / nbd-server.c

diff --git a/nbd-server.c b/nbd-server.c
index 2e3a673..820458c 100644 (file)
--- a/nbd-server.c
+++ b/nbd-server.c
@@ -204,6 +204,7 @@ typedef struct {
        gchar* postrun;      /**< command that will be ran after the client
                                  disconnects */
        gchar* servename;    /**< name of the export as selected by nbd-client */
+       int max_connections; /**< maximum number of opened connections */
 } SERVER;
 
 /**
@@ -358,14 +359,15 @@ inline void writeit(int f, void *buf, size_t len) {
  */
 void usage() {
        printf("This is nbd-server version " VERSION "\n");
-       printf("Usage: [ip:|ip6@]port file_to_export [size][kKmM] [-l authorize_file] [-r] [-m] [-c] [-C configuration file] [-p PID file name] [-o section name]\n"
+       printf("Usage: [ip:|ip6@]port file_to_export [size][kKmM] [-l authorize_file] [-r] [-m] [-c] [-C configuration file] [-p PID file name] [-o section name] [-M max connections]\n"
               "\t-r|--read-only\t\tread only\n"
               "\t-m|--multi-file\t\tmultiple file\n"
               "\t-c|--copy-on-write\tcopy on write\n"
               "\t-C|--config-file\tspecify an alternate configuration file\n"
               "\t-l|--authorize-file\tfile with list of hosts that are allowed to\n\t\t\t\tconnect.\n"
               "\t-p|--pid-file\t\tspecify a filename to write our PID to\n"
-              "\t-o|--output-config\toutput a config file section for what you\n\t\t\t\tspecified on the command line, with the\n\t\t\t\tspecified section name\n\n"
+              "\t-o|--output-config\toutput a config file section for what you\n\t\t\t\tspecified on the command line, with the\n\t\t\t\tspecified section name\n"
+              "\t-M|--max-connections\tspecify the maximum number of opened connections\n\n"
               "\tif port is set to 0, stdin is used (for running from inetd)\n"
               "\tif file_to_export contains '%%s', it is substituted with the IP\n"
               "\t\taddress of the machine trying to connect\n" 
@@ -415,6 +417,7 @@ SERVER* cmdline(int argc, char *argv[]) {
                {"config-file", required_argument, NULL, 'C'},
                {"pid-file", required_argument, NULL, 'p'},
                {"output-config", required_argument, NULL, 'o'},
+               {"max-connection", required_argument, NULL, 'M'},
                {0,0,0,0}
        };
        SERVER *serve;
@@ -431,7 +434,7 @@ SERVER* cmdline(int argc, char *argv[]) {
        serve=g_new0(SERVER, 1);
        serve->authname = g_strdup(default_authname);
        serve->virtstyle=VIRT_IPLIT;
-       while((c=getopt_long(argc, argv, "-C:cl:mo:rp:", long_options, &i))>=0) {
+       while((c=getopt_long(argc, argv, "-C:cl:mo:rp:M:", long_options, &i))>=0) {
                switch (c) {
                case 1:
                        /* non-option argument */
@@ -508,6 +511,9 @@ SERVER* cmdline(int argc, char *argv[]) {
                        g_free(serve->authname);
                        serve->authname=g_strdup(optarg);
                        break;
+               case 'M':
+                       serve->max_connections = strtol(optarg, NULL, 0);
+                       break;
                default:
                        usage();
                        exit(EXIT_FAILURE);
@@ -606,6 +612,8 @@ SERVER* dup_serve(SERVER *s) {
        if(s->servename)
                serve->servename = g_strdup(s->servename);
 
+       serve->max_connections = s->max_connections;
+
        return serve;
 }
 
@@ -705,6 +713,7 @@ GArray* parse_cfile(gchar* f, GError** e) {
                { "sdp",        FALSE,  PARAM_BOOL,     NULL, F_SDP },
                { "sync",       FALSE,  PARAM_BOOL,     NULL, F_SYNC },
                { "listenaddr", FALSE,  PARAM_STRING,   NULL, 0 },
+               { "maxconnections", FALSE, PARAM_INT,   NULL, 0 },
        };
        const int lp_size=sizeof(lp)/sizeof(PARAM);
        PARAM gp[] = {
@@ -755,6 +764,7 @@ GArray* parse_cfile(gchar* f, GError** e) {
                                lp[10].target=lp[11].target=
                                lp[12].target=&(s.flags);
                lp[13].target=&(s.listenaddr);
+               lp[14].target=&(s.max_connections);
 
                /* After the [generic] group, start parsing exports */
                if(i==1) {
@@ -1265,25 +1275,31 @@ CLIENT* negotiate(int net, CLIENT *client, GArray* servers) {
 
                if(!servers)
                        err("programmer error");
-               write(net, &smallflags, sizeof(uint16_t));
-               read(net, &reserved, sizeof(reserved));
-               read(net, &magic, sizeof(magic));
+               if (write(net, &smallflags, sizeof(uint16_t)) < 0)
+                       err("Negotiation failed: %m");
+               if (read(net, &reserved, sizeof(reserved)) < 0)
+                       err("Negotiation failed: %m");
+               if (read(net, &magic, sizeof(magic)) < 0)
+                       err("Negotiation failed: %m");
                magic = ntohll(magic);
                if(magic != opts_magic) {
                        close(net);
                        return NULL;
                }
-               read(net, &opt, sizeof(opt));
+               if (read(net, &opt, sizeof(opt)) < 0)
+                       err("Negotiation failed: %m");
                opt = ntohl(opt);
                if(opt != NBD_OPT_EXPORT_NAME) {
                        close(net);
                        return NULL;
                }
-               read(net, &namelen, sizeof(namelen));
+               if (read(net, &namelen, sizeof(namelen)) < 0)
+                       err("Negotiation failed: %m");
                namelen = ntohl(namelen);
                name = malloc(namelen+1);
                name[namelen]=0;
-               read(net, name, namelen);
+               if (read(net, name, namelen) < 0)
+                       err("Negotiation failed: %m");
                for(i=0; i<servers->len; i++) {
                        SERVER* serve = &(g_array_index(servers, SERVER, i));
                        if(!strcmp(serve->servename, name)) {
@@ -1733,6 +1749,12 @@ int serveloop(GArray* servers) {
                        if(net) {
                                int sock_flags;
 
+                               if(serve->max_connections > 0 &&
+                                  g_hash_table_size(children) >= serve->max_connections) {
+                                       msg2(LOG_INFO, "Max connections reached");
+                                       close(net);
+                                       continue;
+                               }
                                if((sock_flags = fcntl(net, F_GETFL, 0))==-1) {
                                        err("fcntl F_GETFL");
                                }
@@ -2044,7 +2066,7 @@ void glib_message_syslog_redirect(const gchar *log_domain,
       default:
         level=LOG_ERR;
     }
-    syslog(level, message);
+    syslog(level, "%s", message);
 }
 #endif