From e8c3e4c75f9de448ae6e445aa6c10ef33ef27c4f Mon Sep 17 00:00:00 2001
From: yusing <yusing@6uo.me>
Date: Wed, 2 Oct 2024 01:20:25 +0800
Subject: [PATCH] added cidr_whitelist middleware

---
 docs/middlewares.md                           | 28 ++++++-
 .../net/http/middleware/cidr_whitelist.go     | 79 +++++++++++++++++++
 2 files changed, 106 insertions(+), 1 deletion(-)
 create mode 100644 internal/net/http/middleware/cidr_whitelist.go

diff --git a/docs/middlewares.md b/docs/middlewares.md
index 4239e76..319c17f 100644
--- a/docs/middlewares.md
+++ b/docs/middlewares.md
@@ -12,7 +12,7 @@
     - [Real IP](#real-ip)
       - [Custom](#custom)
       - [Cloudflare](#cloudflare)
-    - [Modify request or response](#modify-request-or-response)
+    - [CIDR Whitelist](#cidr-whitelist)
       - [Set headers](#set-headers)
       - [Add headers](#add-headers)
       - [Hide headers](#hide-headers)
@@ -147,6 +147,8 @@ location / {
 }
 ```
 
+[🔼Back to top](#table-of-content)
+
 #### Cloudflare
 
 This is a preset for Cloudflare
@@ -169,6 +171,28 @@ app1:
 
 [🔼Back to top](#table-of-content)
 
+### CIDR Whitelist
+
+```yaml
+# docker labels
+proxy.app1.middlewares.cidr_whitelist.allow: |
+  - 10.0.0.0/8
+  - 192.168.0.0/16
+# optional (default: 403)
+proxy.app1.middlewares.cidr_whitelist.status_code: 403
+# optional (default: "IP not allowed")
+proxy.app1.middlewares.cidr_whitelist.message: "IP not allowed"
+
+# include file
+app1:
+  middlewares:
+    whitelist:
+      cidr:
+        - 10.0.0.0/8
+        - 192.168.0.0/16
+      status_code: 403          # default
+      message: "IP not allowed" # default
+
 ### Modify request or response
 
 ```yaml
@@ -365,6 +389,8 @@ You may use them with `<middleware_name>@file`
 
 See [example](../internal/net/http/middleware/test_data/middleware_compose.yml)
 
+[🔼Back to top](#table-of-content)
+
 ## Examples
 
 ### Authentik (untested, experimental)
diff --git a/internal/net/http/middleware/cidr_whitelist.go b/internal/net/http/middleware/cidr_whitelist.go
new file mode 100644
index 0000000..cf2c62b
--- /dev/null
+++ b/internal/net/http/middleware/cidr_whitelist.go
@@ -0,0 +1,79 @@
+package middleware
+
+import (
+	"net"
+	"net/http"
+
+	D "github.com/yusing/go-proxy/internal/docker"
+	E "github.com/yusing/go-proxy/internal/error"
+	"github.com/yusing/go-proxy/internal/types"
+)
+
+type cidrWhitelist struct {
+	*cidrWhitelistOpts
+	m *Middleware
+}
+
+type cidrWhitelistOpts struct {
+	Allow      []*types.CIDR
+	StatusCode int
+	Message    string
+
+	trustedAddr map[string]struct{} // cache for trusted IPs
+}
+
+var CIDRWhiteList = &cidrWhitelist{
+	m: &Middleware{
+		labelParserMap: D.ValueParserMap{
+			"allow":      D.YamlStringListParser,
+			"statusCode": D.IntParser,
+		},
+	},
+}
+
+var cidrWhitelistDefaults = func() *cidrWhitelistOpts {
+	return &cidrWhitelistOpts{
+		Allow:       []*types.CIDR{},
+		StatusCode:  http.StatusForbidden,
+		Message:     "IP not allowed",
+		trustedAddr: make(map[string]struct{}),
+	}
+}
+
+func NewCIDRWhitelist(opts OptionsRaw) (*Middleware, E.NestedError) {
+	wl := new(cidrWhitelist)
+	wl.m = &Middleware{
+		impl:   wl,
+		before: wl.checkIP,
+	}
+	wl.cidrWhitelistOpts = cidrWhitelistDefaults()
+	err := Deserialize(opts, wl.cidrWhitelistOpts)
+	if err != nil {
+		return nil, err
+	}
+	if len(wl.cidrWhitelistOpts.Allow) == 0 {
+		return nil, E.Missing("allow range")
+	}
+	return wl.m, nil
+}
+
+func (wl *cidrWhitelist) checkIP(next http.Handler, w ResponseWriter, r *Request) {
+	var ok bool
+	if _, ok = wl.trustedAddr[r.RemoteAddr]; !ok {
+		ip := net.IP(r.RemoteAddr)
+		for _, cidr := range wl.cidrWhitelistOpts.Allow {
+			if cidr.Contains(ip) {
+				wl.trustedAddr[r.RemoteAddr] = struct{}{}
+				ok = true
+				break
+			}
+		}
+	}
+	if !ok {
+		w.WriteHeader(wl.StatusCode)
+		w.Write([]byte(wl.Message))
+		return
+	}
+
+	next.ServeHTTP(w, r)
+}