From ef5d0207f13c5226283421dc3b5370e09dc4fa4d Mon Sep 17 00:00:00 2001
From: Michael Khalturin <halturinmv@gmail.com>
Date: Tue, 27 Aug 2024 09:48:25 +0300
Subject: [PATCH] added allow_ping config option

---
 config/README.md          |  3 +++
 config/config.go          |  3 +++
 config/examples/debug.yml | 52 +++++++++++++++++++++++++++++++++++++++
 main.go                   | 33 ++++++++++++++++++++++++-
 4 files changed, 90 insertions(+), 1 deletion(-)
 create mode 100644 config/examples/debug.yml

diff --git a/config/README.md b/config/README.md
index 49bc85cc..bb2f27b2 100644
--- a/config/README.md
+++ b/config/README.md
@@ -16,6 +16,9 @@ log_debug: <bool> | default = false [optional]
 # Whether to ignore security warnings
 hack_me_please: <bool> | default = false [optional]
 
+# Allow ping server
+allow_ping: <bool> | default = false [optional]
+
 # Named list of cache configurations
 caches:
   - <cache_config> ...
diff --git a/config/config.go b/config/config.go
index 7d32d08d..652cfe9b 100644
--- a/config/config.go
+++ b/config/config.go
@@ -78,6 +78,9 @@ type Config struct {
 
 	ConnectionPool ConnectionPool `yaml:"connection_pool,omitempty"`
 
+	// Allow to proxy ping requests
+	AllowPing bool `yaml:"allow_ping,omitempty"`
+
 	networkReg map[string]Networks
 
 	// Catches all undefined fields
diff --git a/config/examples/debug.yml b/config/examples/debug.yml
new file mode 100644
index 00000000..cb88dd05
--- /dev/null
+++ b/config/examples/debug.yml
@@ -0,0 +1,52 @@
+log_debug: true
+hack_me_please: true
+allow_ping: true
+
+caches:
+  - name: "longterm"
+    mode: "file_system"
+    file_system:
+      dir: "/tmp/chproxy/longterm/cache"
+      max_size: 10Gb
+    expire: 1h
+    shared_with_all_users: true
+
+max_error_reason_size: 10GB
+
+server:
+  http:
+    listen_addr: ":80"
+    idle_timeout: 20m
+
+  https:
+    listen_addr: ":4433"
+    autocert:
+      cache_dir: "/tmp/chproxy/certs"
+  proxy:
+    enable: true
+    header: CF-Connecting-IP
+
+users:
+  - name: "*"
+    to_cluster: "dl"
+    to_user: "*"
+    is_wildcarded: true
+    max_concurrent_queries: 4
+    max_execution_time: 15m
+    deny_https: false
+    cache: "longterm"
+
+clusters:
+  - name: "dl"
+    scheme: "http"
+    nodes: ["localhost:28123"]
+    # heartbeat:
+    #   interval: 1m
+    #   timeout: 10s
+    #   request: "/?query=SELECT%201%2B1"
+    # q  response: "2\n"
+    kill_query_user:
+      name: "chproxy"
+      password: "chproxy"
+    users:
+      - name: "*"
diff --git a/main.go b/main.go
index b24da71e..ae19ce03 100644
--- a/main.go
+++ b/main.go
@@ -35,6 +35,7 @@ var (
 	allowedNetworksHTTPS   atomic.Value
 	allowedNetworksMetrics atomic.Value
 	proxyHandler           atomic.Value
+	allowPing              atomic.Bool
 )
 
 func main() {
@@ -237,12 +238,41 @@ func serveHTTP(rw http.ResponseWriter, r *http.Request) {
 		}
 		proxy.refreshCacheMetrics()
 		promHandler.ServeHTTP(rw, r)
-	case "/", "/query", "/ping":
+	case "/", "/query":
 		var err error
 		// nolint:forcetypeassert // We will cover this by tests as we control what is stored.
 		proxyHandler := proxyHandler.Load().(*ProxyHandler)
 		r.RemoteAddr = proxyHandler.GetRemoteAddr(r)
 
+		var an *config.Networks
+		if r.TLS != nil {
+			// nolint:forcetypeassert // We will cover this by tests as we control what is stored.
+			an = allowedNetworksHTTPS.Load().(*config.Networks)
+			err = fmt.Errorf("https connections are not allowed from %s", r.RemoteAddr)
+		} else {
+			// nolint:forcetypeassert // We will cover this by tests as we control what is stored.
+			an = allowedNetworksHTTP.Load().(*config.Networks)
+			err = fmt.Errorf("http connections are not allowed from %s", r.RemoteAddr)
+		}
+		if !an.Contains(r.RemoteAddr) {
+			rw.Header().Set("Connection", "close")
+			respondWith(rw, err, http.StatusForbidden)
+			return
+		}
+		proxy.ServeHTTP(rw, r)
+	case "/ping":
+		var err error
+
+		if !allowPing.Load() {
+			err = fmt.Errorf("ping is not allowed")
+			respondWith(rw, err, http.StatusForbidden)
+			return
+		}
+
+		// nolint:forcetypeassert // We will cover this by tests as we control what is stored.
+		proxyHandler := proxyHandler.Load().(*ProxyHandler)
+		r.RemoteAddr = proxyHandler.GetRemoteAddr(r)
+
 		var an *config.Networks
 		if r.TLS != nil {
 			// nolint:forcetypeassert // We will cover this by tests as we control what is stored.
@@ -296,6 +326,7 @@ func applyConfig(cfg *config.Config) error {
 	allowedNetworksHTTPS.Store(&cfg.Server.HTTPS.AllowedNetworks)
 	allowedNetworksMetrics.Store(&cfg.Server.Metrics.AllowedNetworks)
 	proxyHandler.Store(NewProxyHandler(&cfg.Server.Proxy))
+	allowPing.Store(cfg.AllowPing)
 	log.SetDebug(cfg.LogDebug)
 	log.Infof("Loaded config:\n%s", cfg)