From 0c3cb5e6aa4cc0ea447f2be285684890fe208d35 Mon Sep 17 00:00:00 2001
From: LittleChest <hiya@littlew.top>
Date: Wed, 13 Aug 2025 09:14:41 +0800
Subject: [PATCH] handler: Custom prefix length

---
 README.md                            | 10 ++++++----
 _worker.js                           |  9 ++++++++-
 common.js                            | 14 ++++++++------
 middleware.js                        |  8 +++++++-
 netlify/edge-functions/middleware.js |  9 ++++++++-
 5 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/README.md b/README.md
index 415874e..09c4670 100644
--- a/README.md
+++ b/README.md
@@ -19,10 +19,12 @@ Read [Dohna NS Documentation](https://dohna.ovh/) to learn how to install Dohna
 
 ## Environment Variables
 
-| Key | Default                      | Description                                      |
-| --- | ---------------------------- | ------------------------------------------------ |
-| DNS | https://dns.google/dns-query | Specify a DNS over HTTPS server as the upstream. |
-| API | https://dns.google/resolve   | Specify a JSON API server as the upstream.       |
+| Key         | Default                      | Description                                        |
+| ----------- | ---------------------------- | -------------------------------------------------- |
+| DNS         | https://dns.google/dns-query | Specify a DNS over HTTPS server as the upstream.   |
+| API         | https://dns.google/resolve   | Specify a JSON API server as the upstream.         |
+| IPV4_PREFIX | 32                           | Specify the EDNS client subnet IPv4 prefix length. |
+| IPV6_PREFIX | 128                          | Specify the EDNS client subnet IPv6 prefix length. |
 
 ## Self-hosted
 
diff --git a/_worker.js b/_worker.js
index 55a9a12..508988c 100644
--- a/_worker.js
+++ b/_worker.js
@@ -2,5 +2,12 @@ import handler from "./common";
 
 export default {
   fetch: async (request, env) =>
-    handler(request, env.DNS, env.API, request.headers.get("cf-connecting-ip")),
+    handler(
+      request,
+      env.DNS,
+      env.API,
+      env.IPV4_PREFIX,
+      env.IPV6_PREFIX,
+      request.headers.get("cf-connecting-ip")
+    ),
 };
diff --git a/common.js b/common.js
index a006b3d..2fe55ad 100644
--- a/common.js
+++ b/common.js
@@ -2,6 +2,8 @@ export default async function handler(
   request,
   dns = "https://dns.google/dns-query",
   api = "https://dns.google/resolve",
+  ipv4Prefix = 32,
+  ipv6Prefix = 128,
   rawIP
 ) {
   const { method, headers, url } = request;
@@ -75,14 +77,14 @@ export default async function handler(
     }
 
     if (queryData !== undefined) {
-      res = await queryDns(queryData, ip, dns);
+      res = await queryDns(queryData, ip, dns, ipv4Prefix, ipv6Prefix);
     }
   }
 
   return res;
 }
 
-async function queryDns(queryData, ip, dns) {
+async function queryDns(queryData, ip, dns, ipv4Prefix, ipv6Prefix) {
   const hasOptRecord = checkForOptRecord(queryData);
   let newQueryData = queryData;
   if (!hasOptRecord && ip) {
@@ -90,7 +92,7 @@ async function queryDns(queryData, ip, dns) {
     const [headerAndQuestion] = extractHeaderAndQuestion(queryData);
 
     // Construct a new OPT record with ECS option
-    const optRecord = createOptRecord(ip);
+    const optRecord = createOptRecord(ip, ipv4Prefix, ipv6Prefix);
 
     // Combine the header, question, and new OPT record to create a new query
     newQueryData = combineQueryData(headerAndQuestion, optRecord);
@@ -172,21 +174,21 @@ function extractHeaderAndQuestion(data) {
   return [headerAndQuestion, offset];
 }
 
-function createOptRecord(ip) {
+function createOptRecord(ip, ipv4Prefix, ipv6Prefix) {
   let ecsData;
   let family;
 
   if (isIPv4(ip)) {
+    const prefixLength = ipv4Prefix
     // Convert client IP to bytes
     const ipParts = ip.split(".").map((part) => parseInt(part, 10));
     family = 1; // IPv4
-    const prefixLength = 32; // Adjust the prefix length as needed
     ecsData = [0, 8, 0, 8, 0, family, prefixLength, 0, ...ipParts];
   } else if (isIPv6(ip)) {
+    const prefixLength = ipv6Prefix
     // Convert client IP to bytes
     const ipParts = ipv6ToBytes(ip);
     family = 2; // IPv6
-    const prefixLength = 128; // Adjust the prefix length as needed
     ecsData = [0, 8, 0, 20, 0, family, prefixLength, 0, ...ipParts];
   } else {
     throw new Error("Invalid IP address");
diff --git a/middleware.js b/middleware.js
index 2259bb7..ad4684b 100644
--- a/middleware.js
+++ b/middleware.js
@@ -1,5 +1,11 @@
 import handler from "./common";
 
 export default middleware = async (request) => {
-  return handler(request, process.env.DNS, process.env.API);
+  return handler(
+    request,
+    process.env.DNS,
+    process.env.API,
+    process.env.IPV4_PREFIX,
+    process.env.IPV6_PREFIX
+  );
 };
diff --git a/netlify/edge-functions/middleware.js b/netlify/edge-functions/middleware.js
index 5c16457..3062f7f 100644
--- a/netlify/edge-functions/middleware.js
+++ b/netlify/edge-functions/middleware.js
@@ -1,4 +1,11 @@
 import handler from "../../common.js";
 export default async (request) =>
-  handler(request, Netlify.env.get("DNS"), Netlify.env.get("API"), Netlify.context.ip);
+  handler(
+    request,
+    Netlify.env.get("DNS"),
+    Netlify.env.get("API"),
+    Netlify.env.get("IPV4_PREFIX"),
+    Netlify.env.get("IPV6_PREFIX"),
+    Netlify.context.ip
+  );
 export const config = { path: "*" };