Limit simultaneous clients amount

[gocheese.git] / gocheese.go

diff --git a/gocheese.go b/gocheese.go
index 835e8c5b392a2f66ab0c2a7e0500c0a976445477..26c2c27ba74673cda10e03eb6b44e31ee924920f 100644 (file)
--- a/gocheese.go
+++ b/gocheese.go
@@ -1,6 +1,7 @@
 /*
 GoCheese -- Python private package repository and caching proxy
 Copyright (C) 2019 Sergey Matveev <stargrave@stargrave.org>
+              2019 Elena Balakhonova <balakhonova_e@riseup.net>
 
 This program is free software: you can redistribute it and/or modify
 it under the terms of the GNU General Public License as published by
@@ -20,6 +21,7 @@ package main
 
 import (
        "bytes"
+       "context"
        "crypto/sha256"
        "encoding/hex"
        "flag"
@@ -27,6 +29,7 @@ import (
        "io"
        "io/ioutil"
        "log"
+       "net"
        "net/http"
        "net/url"
        "os"
@@ -36,6 +39,9 @@ import (
        "runtime"
        "strings"
        "syscall"
+       "time"
+
+       "golang.org/x/net/netutil"
 )
 
 const (
@@ -68,6 +74,7 @@ var (
        passwdPath       = flag.String("passwd", "passwd", "Path to file with authenticators")
        passwdCheck      = flag.Bool("passwd-check", false, "Test the -passwd file for syntax errors and exit")
        fsck             = flag.Bool("fsck", false, "Check integrity of all packages")
+       maxClients       = flag.Int("maxclients", 128, "Maximal amount of simultaneous clients")
        version          = flag.Bool("version", false, "Print version information")
        warranty         = flag.Bool("warranty", false, "Print warranty information")
 
@@ -441,14 +448,37 @@ func main() {
        refreshPasswd()
        log.Println("root:", *root, "bind:", *bind)
        needsRefreshPasswd := make(chan os.Signal, 0)
+       needsShutdown := make(chan os.Signal, 0)
+       killed := make(chan error, 0)
+       http.HandleFunc(*norefreshURLPath, handler)
+       http.HandleFunc(*refreshURLPath, handler)
+       ln, err := net.Listen("tcp", *bind)
+       if err != nil {
+               log.Fatal(err)
+       }
+       s := &http.Server{
+               ReadTimeout:  time.Minute,
+               WriteTimeout: time.Minute,
+       }
        signal.Notify(needsRefreshPasswd, syscall.SIGHUP)
+       signal.Notify(needsShutdown, syscall.SIGTERM, syscall.SIGINT)
        go func() {
                for range needsRefreshPasswd {
                        log.Println("Refreshing passwords")
                        refreshPasswd()
                }
        }()
-       http.HandleFunc(*norefreshURLPath, handler)
-       http.HandleFunc(*refreshURLPath, handler)
-       log.Fatal(http.ListenAndServe(*bind, nil))
+       go func(s *http.Server) {
+               <-needsShutdown
+               log.Println("Shutting down")
+               ctx, cancel := context.WithTimeout(context.TODO(), time.Minute)
+               killed <- s.Shutdown(ctx)
+               cancel()
+       }(s)
+       if err := s.Serve(netutil.LimitListener(ln, *maxClients)); err != http.ErrServerClosed {
+               log.Fatal(err)
+       }
+       if err := <-killed; err != nil {
+               log.Fatal(err)
+       }
 }