From 06c363e55a60bb0d063e206c457615e77d0b8ed1 Mon Sep 17 00:00:00 2001
From: WeebDataHoarder <weebdatahoarder@noreply.gammaspectra.live>
Date: Sun, 27 Apr 2025 17:37:34 +0200
Subject: [PATCH] context: add ip prefix on keyed cookie

---
 lib/challenge/data.go | 17 ++++++++++++++++-
 lib/challenge/key.go  | 14 ++------------
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/lib/challenge/data.go b/lib/challenge/data.go
index 7adac9e..b4c1841 100644
--- a/lib/challenge/data.go
+++ b/lib/challenge/data.go
@@ -91,9 +91,11 @@ func CreateRequestData(r *http.Request, state StateInterface) (*http.Request, *R
 	sum := sha256.New()
 	sum.Write([]byte(r.Host))
 	sum.Write([]byte{0})
+	sum.Write(data.NetworkPrefix().AsSlice())
+	sum.Write([]byte{0})
 	sum.Write(state.PublicKey())
 	sum.Write([]byte{0})
-	data.CookiePrefix = utils.CookiePrefix + hex.EncodeToString(sum.Sum(nil)[:4]) + "-"
+	data.CookiePrefix = utils.CookiePrefix + hex.EncodeToString(sum.Sum(nil)[:6]) + "-"
 
 	r = r.WithContext(context.WithValue(r.Context(), requestDataContextKey{}, &data))
 	r = utils.SetRemoteAddress(r, data.RemoteAddress)
@@ -129,6 +131,19 @@ func (d *RequestData) Parent() cel.Activation {
 	return nil
 }
 
+func (d *RequestData) NetworkPrefix() netip.Addr {
+	address := d.RemoteAddress.Addr().Unmap()
+	if address.Is4() {
+		// Take a /24 for IPv4
+		prefix, _ := address.Prefix(24)
+		return prefix.Addr()
+	} else {
+		// Take a /64 for IPv6
+		prefix, _ := address.Prefix(64)
+		return prefix.Addr()
+	}
+}
+
 func (d *RequestData) SetOpt(n, v string) {
 	d.opts[n] = v
 }
diff --git a/lib/challenge/key.go b/lib/challenge/key.go
index 22466c4..7b15932 100644
--- a/lib/challenge/key.go
+++ b/lib/challenge/key.go
@@ -42,22 +42,12 @@ func KeyFromString(s string) (Key, error) {
 
 func GetChallengeKeyForRequest(state StateInterface, reg *Registration, until time.Time, r *http.Request) Key {
 	data := RequestDataFromContext(r.Context())
-	address := data.RemoteAddress.Addr().Unmap()
-	var keyAddr [16]byte
-	if address.Is4() {
-		// Take a /24 for IPv4
-		prefix, _ := address.Prefix(24)
-		keyAddr = prefix.Addr().As16()
-	} else {
-		// Take a /64 for IPv6
-		prefix, _ := address.Prefix(64)
-		keyAddr = prefix.Addr().As16()
-	}
 
 	hasher := sha256.New()
 	hasher.Write([]byte("challenge\x00"))
 	hasher.Write([]byte(reg.Name))
 	hasher.Write([]byte{0})
+	keyAddr := data.NetworkPrefix().As16()
 	hasher.Write(keyAddr[:])
 	hasher.Write([]byte{0})
 
@@ -83,7 +73,7 @@ func GetChallengeKeyForRequest(state StateInterface, reg *Registration, until ti
 
 	sum[0] = 0
 
-	if address.Is4() {
+	if data.RemoteAddress.Addr().Unmap().Is4() {
 		// Is IPv4, mark
 		sum.Set(KeyFlagIsIPv4)
 	}