filters/limit/limit.c

   1 /* nbdkit
   2  * Copyright (C) 2019-2020 Red Hat Inc.
   3  *
   4  * Redistribution and use in source and binary forms, with or without
   5  * modification, are permitted provided that the following conditions are
   6  * met:
   7  *
   8  * * Redistributions of source code must retain the above copyright
   9  * notice, this list of conditions and the following disclaimer.
  10  *
  11  * * Redistributions in binary form must reproduce the above copyright
  12  * notice, this list of conditions and the following disclaimer in the
  13  * documentation and/or other materials provided with the distribution.
  14  *
  15  * * Neither the name of Red Hat nor the names of its contributors may be
  16  * used to endorse or promote products derived from this software without
  17  * specific prior written permission.
  18  *
  19  * THIS SOFTWARE IS PROVIDED BY RED HAT AND CONTRIBUTORS ''AS IS'' AND
  20  * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
  21  * THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
  22  * PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL RED HAT OR
  23  * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  24  * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  25  * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
  26  * USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
  27  * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
  28  * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
  29  * OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
  30  * SUCH DAMAGE.
  31  */
  32
  33 #include <config.h>
  34
  35 #include <stdio.h>
  36 #include <stdlib.h>
  37 #include <string.h>
  38 #include <pthread.h>
  39
  40 #include <nbdkit-filter.h>
  41
  42 #include "cleanup.h"
  43
  44 /* Counts client connections. */
  45 static unsigned connections;
  46 static pthread_mutex_t lock = PTHREAD_MUTEX_INITIALIZER;
  47
  48 /* Client limit (0 = filter is disabled). */
  49 static unsigned limit = 1;
  50
  51 static int
  52 limit_config (nbdkit_next_config *next, nbdkit_backend *nxdata,
  53               const char *key, const char *value)
  54 {
  55   if (strcmp (key, "limit") == 0) {
  56     if (nbdkit_parse_unsigned ("limit", value, &limit) == -1)
  57       return -1;
  58     return 0;
  59   }
  60
  61   return next (nxdata, key, value);
  62 }
  63
  64 static void
  65 too_many_clients_error (void)
  66 {
  67   nbdkit_error ("limit: too many clients connected, connection rejected");
  68 }
  69
  70 /* We limit connections in the preconnect stage (in particular before
  71  * any heavyweight NBD or TLS negotiations has been done).  However we
  72  * count connections in the open/close calls since clients can drop
  73  * out between preconnect and open.
  74  */
  75 static int
  76 limit_preconnect (nbdkit_next_preconnect *next, nbdkit_backend *nxdata,
  77                   int readonly)
  78 {
  79   if (next (nxdata, readonly) == -1)
  80     return -1;
  81
  82   ACQUIRE_LOCK_FOR_CURRENT_SCOPE (&lock);
  83
  84   if (limit > 0 && connections >= limit) {
  85     too_many_clients_error ();
  86     return -1;
  87   }
  88
  89   return 0;
  90 }
  91
  92 static void *
  93 limit_open (nbdkit_next_open *next, nbdkit_context *nxdata,
  94             int readonly, const char *exportname, int is_tls)
  95 {
  96   if (next (nxdata, readonly, exportname) == -1)
  97     return NULL;
  98
  99   ACQUIRE_LOCK_FOR_CURRENT_SCOPE (&lock);
 100
 101   /* We have to check again because clients can artificially slow down
 102    * the NBD negotiation in order to bypass the limit otherwise.
 103    */
 104   if (limit > 0 && connections >= limit) {
 105     too_many_clients_error ();
 106     return NULL;
 107   }
 108
 109   connections++;
 110   return NBDKIT_HANDLE_NOT_NEEDED;
 111 }
 112
 113 static void
 114 limit_close (void *handle)
 115 {
 116   ACQUIRE_LOCK_FOR_CURRENT_SCOPE (&lock);
 117   connections--;
 118 }
 119
 120 static struct nbdkit_filter filter = {
 121   .name              = "limit",
 122   .longname          = "nbdkit limit filter",
 123   .config            = limit_config,
 124   .preconnect        = limit_preconnect,
 125   .open              = limit_open,
 126   .close             = limit_close,
 127 };
 128
 129 NBDKIT_REGISTER_FILTER (filter)