From ab15dc4bb5f530a6459032cc49922b45f569b4cb Mon Sep 17 00:00:00 2001
From: WeebDataHoarder <57538841+WeebDataHoarder@users.noreply.github.com>
Date: Fri, 5 Aug 2022 00:26:22 +0200
Subject: [PATCH] Use server preference on cipher suites and curves

---
 go.mod          |  1 +
 go.sum          |  2 ++
 tlsutils/tls.go | 68 ++++++++++++++++++++++++++++++++++++++++++-------
 3 files changed, 62 insertions(+), 9 deletions(-)

diff --git a/go.mod b/go.mod
index 2a0b90e..d5630ee 100644
--- a/go.mod
+++ b/go.mod
@@ -9,6 +9,7 @@ require (
 	github.com/lucas-clemente/quic-go v0.28.1
 	github.com/multiformats/go-multihash v0.2.0
 	github.com/valyala/fasthttp v1.38.0
+	golang.org/x/exp v0.0.0-20220722155223-a9213eeb770e
 )
 
 require (
diff --git a/go.sum b/go.sum
index 8635eb9..6f1cabc 100644
--- a/go.sum
+++ b/go.sum
@@ -213,6 +213,8 @@ golang.org/x/crypto v0.0.0-20220315160706-3147a52a75dd/go.mod h1:IxCIyHEi3zRg3s0
 golang.org/x/crypto v0.0.0-20220722155217-630584e8d5aa h1:zuSxTR4o9y82ebqCUJYNGJbGPo6sKVl54f/TVDObg1c=
 golang.org/x/crypto v0.0.0-20220722155217-630584e8d5aa/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
+golang.org/x/exp v0.0.0-20220722155223-a9213eeb770e h1:+WEEuIdZHnUeJJmEUjyYC2gfUMj69yZXw17EnHg/otA=
+golang.org/x/exp v0.0.0-20220722155223-a9213eeb770e/go.mod h1:Kr81I6Kryrl9sr8s2FK3vxD90NdsKWRuOIl2O4CvYbA=
 golang.org/x/lint v0.0.0-20180702182130-06c8688daad7/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
 golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
 golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU=
diff --git a/tlsutils/tls.go b/tlsutils/tls.go
index 96b682b..fb292d4 100644
--- a/tlsutils/tls.go
+++ b/tlsutils/tls.go
@@ -8,6 +8,7 @@ import (
 	"crypto/x509"
 	"crypto/x509/pkix"
 	"encoding/pem"
+	"golang.org/x/exp/slices"
 	"math"
 	"math/big"
 	"time"
@@ -91,6 +92,9 @@ func NewTLSConfiguration(certificatePath, keypairPath, sni string) (*Configurati
 			tls.CurveP384,
 		},
 		CipherSuites: []uint16{
+			tls.TLS_CHACHA20_POLY1305_SHA256,
+			tls.TLS_AES_256_GCM_SHA384,
+			tls.TLS_AES_128_GCM_SHA256,
 			tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256,
 			tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
 			tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
@@ -115,16 +119,62 @@ func NewTLSConfiguration(certificatePath, keypairPath, sni string) (*Configurati
 		}...)
 	}
 
-	return &Configuration{
-		Config: tlsConfig,
-		QUICConfig: &tls.Config{
-			GetCertificate: func(info *tls.ClientHelloInfo) (*tls.Certificate, error) {
-				if len(sni) == 0 || sni == info.ServerName {
-					return &serverCertificate, nil
-				}
-				return bogusCertificate, nil
-			},
+	tlsConfig.GetConfigForClient = func(info *tls.ClientHelloInfo) (*tls.Config, error) {
+		configClone := tlsConfig.Clone()
+
+		//Have proper server preference
+		for _, suite := range configClone.CipherSuites {
+			if slices.Contains(info.CipherSuites, suite) {
+				configClone.CipherSuites = []uint16{suite}
+				break
+			}
+		}
+
+		//Have proper server preference
+		for _, curve := range configClone.CurvePreferences {
+			if slices.Contains(info.SupportedCurves, curve) {
+				configClone.CurvePreferences = []tls.CurveID{curve}
+				break
+			}
+		}
+
+		return configClone, nil
+	}
+
+	quicTlsConfig := &tls.Config{
+		GetCertificate: func(info *tls.ClientHelloInfo) (*tls.Certificate, error) {
+			if len(sni) == 0 || sni == info.ServerName {
+				return &serverCertificate, nil
+			}
+			return bogusCertificate, nil
 		},
+	}
+
+	quicTlsConfig.GetConfigForClient = func(info *tls.ClientHelloInfo) (*tls.Config, error) {
+		configClone := quicTlsConfig.Clone()
+
+		//Have proper server suite preference
+		for _, suite := range configClone.CipherSuites {
+			if slices.Contains(info.CipherSuites, suite) {
+				configClone.CipherSuites = []uint16{suite}
+				break
+			}
+		}
+
+		//Have proper server curve preference
+		for _, curve := range configClone.CurvePreferences {
+			if slices.Contains(info.SupportedCurves, curve) {
+				configClone.CurvePreferences = []tls.CurveID{curve}
+				break
+			}
+		}
+
+		return configClone, nil
+	}
+
+	return &Configuration{
+		Config:     tlsConfig,
+		QUICConfig: quicTlsConfig,
 	}, nil
 }