From a0b3e049d03056bd2c8bc8163130f5a9692df5c9 Mon Sep 17 00:00:00 2001
From: Kane York <rikingcoding@gmail.com>
Date: Mon, 13 Nov 2017 12:48:56 -0800
Subject: [PATCH] Reject over capacity before finishing TLS handshake

---
 .../cmd/ffzsocketserver/socketserver.go       |  3 +-
 socketserver/server/handlecore.go             | 37 +++++++++++++------
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/socketserver/cmd/ffzsocketserver/socketserver.go b/socketserver/cmd/ffzsocketserver/socketserver.go
index 1f7f2db0..cea1bb07 100644
--- a/socketserver/cmd/ffzsocketserver/socketserver.go
+++ b/socketserver/cmd/ffzsocketserver/socketserver.go
@@ -81,7 +81,8 @@ func main() {
 			Addr:    conf.SSLListenAddr,
 			Handler: http.DefaultServeMux,
 			TLSConfig: &tls.Config{
-				GetCertificate: reloader.GetCertificateFunc(),
+				GetCertificate:     reloader.GetCertificateFunc(),
+				GetConfigForClient: server.TLSEarlyReject,
 			},
 		}
 		go func() {
diff --git a/socketserver/server/handlecore.go b/socketserver/server/handlecore.go
index fa3158fe..a1383a91 100644
--- a/socketserver/server/handlecore.go
+++ b/socketserver/server/handlecore.go
@@ -2,6 +2,7 @@ package server // import "github.com/FrankerFaceZ/FrankerFaceZ/socketserver/serv
 
 import (
 	"bytes"
+	"crypto/tls"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -239,6 +240,29 @@ var BannerHTML []byte
 // StopAcceptingConnectionsCh is closed while the server is shutting down.
 var StopAcceptingConnectionsCh = make(chan struct{})
 
+func shouldRejectConnection() bool {
+	memFreeKB := atomic.LoadUint64(&Statistics.SysMemFreeKB)
+	if memFreeKB > 0 && memFreeKB < Configuration.MinMemoryKBytes {
+		return true
+	}
+
+	curClients := atomic.LoadUint64(&Statistics.CurrentClientCount)
+	if Configuration.MaxClientCount != 0 && curClients >= Configuration.MaxClientCount {
+		return true
+	}
+
+	return false
+}
+
+var errEarlyTLSReject = errors.New("over capacity")
+
+func TLSEarlyReject(*tls.ClientHelloInfo) (*tls.Config, error) {
+	if shouldRejectConnection() {
+		return nil, errEarlyTLSReject
+	}
+	return nil, nil
+}
+
 // HTTPHandleRootURL is the http.HandleFunc for requests on `/`.
 // It either uses the SocketUpgrader or writes out the BannerHTML.
 func HTTPHandleRootURL(w http.ResponseWriter, r *http.Request) {
@@ -251,21 +275,12 @@ func HTTPHandleRootURL(w http.ResponseWriter, r *http.Request) {
 	if strings.Contains(strings.ToLower(r.Header.Get("Connection")), "upgrade") {
 		updateSysMem()
 
-		if Statistics.SysMemFreeKB > 0 && Statistics.SysMemFreeKB < Configuration.MinMemoryKBytes {
+		if shouldRejectConnection() {
 			w.WriteHeader(503)
-			fmt.Fprint(w, "error: low memory")
+			fmt.Fprint(w, "connection rejected: over capacity")
 			return
 		}
 
-		if Configuration.MaxClientCount != 0 {
-			curClients := atomic.LoadUint64(&Statistics.CurrentClientCount)
-			if curClients >= Configuration.MaxClientCount {
-				w.WriteHeader(503)
-				fmt.Fprint(w, "error: client limit reached")
-				return
-			}
-		}
-
 		conn, err := SocketUpgrader.Upgrade(w, r, nil)
 		if err != nil {
 			fmt.Fprintf(w, "error: %v", err)