From 94383f29e919f32ae25c6bcca9270349fb31c719 Mon Sep 17 00:00:00 2001
From: mleku <me@mleku.dev>
Date: Sun, 31 Aug 2025 10:00:50 +0100
Subject: [PATCH] Add IP whitelist configuration, enhance message handling with
 envelope identification, and log remote addresses for improved connection
 control

---
 app/config/config.go    | 15 ++++++++-----
 app/handle-message.go   | 50 +++++++++++++++++++++++++++++++++++++++--
 app/handle-websocket.go | 16 +++++++++++--
 3 files changed, 71 insertions(+), 10 deletions(-)

diff --git a/app/config/config.go b/app/config/config.go
index 66f0bf9..65cc93c 100644
--- a/app/config/config.go
+++ b/app/config/config.go
@@ -16,6 +16,7 @@ import (
 	"go-simpler.org/env"
 	lol "lol.mleku.dev"
 	"lol.mleku.dev/chk"
+	"lol.mleku.dev/log"
 	"next.orly.dev/pkg/version"
 )
 
@@ -23,12 +24,13 @@ import (
 // and default values. It defines parameters for app behaviour, storage
 // locations, logging, and network settings used across the relay service.
 type C struct {
-	AppName  string `env:"ORLY_APP_NAME" usage:"set a name to display on information about the relay" default:"ORLY"`
-	DataDir  string `env:"ORLY_DATA_DIR" usage:"storage location for the event store" default:"~/.local/share/ORLY"`
-	Listen   string `env:"ORLY_LISTEN" default:"0.0.0.0" usage:"network listen address"`
-	Port     int    `env:"ORLY_PORT" default:"3334" usage:"port to listen on"`
-	LogLevel string `env:"ORLY_LOG_LEVEL" default:"info" usage:"debug level: fatal error warn info debug trace"`
-	Pprof    string `env:"ORLY_PPROF" usage:"enable pprof in modes: cpu,memory,allocation"`
+	AppName     string   `env:"ORLY_APP_NAME" usage:"set a name to display on information about the relay" default:"ORLY"`
+	DataDir     string   `env:"ORLY_DATA_DIR" usage:"storage location for the event store" default:"~/.local/share/ORLY"`
+	Listen      string   `env:"ORLY_LISTEN" default:"0.0.0.0" usage:"network listen address"`
+	Port        int      `env:"ORLY_PORT" default:"3334" usage:"port to listen on"`
+	LogLevel    string   `env:"ORLY_LOG_LEVEL" default:"info" usage:"debug level: fatal error warn info debug trace"`
+	Pprof       string   `env:"ORLY_PPROF" usage:"enable pprof in modes: cpu,memory,allocation"`
+	IPWhitelist []string `env:"ORLY_IP_WHITELIST" usage:"comma-separated list of IP addresses to allow access from, matches on prefixes to allow private subnets, eg 10.0.0 = 10.0.0.0/8"`
 }
 
 // New creates and initializes a new configuration object for the relay
@@ -69,6 +71,7 @@ func New() (cfg *C, err error) {
 		os.Exit(0)
 	}
 	lol.SetLogLevel(cfg.LogLevel)
+	log.I.S(cfg.IPWhitelist)
 	return
 }
 
diff --git a/app/handle-message.go b/app/handle-message.go
index 325b323..82892e8 100644
--- a/app/handle-message.go
+++ b/app/handle-message.go
@@ -1,9 +1,55 @@
 package app
 
 import (
+	"fmt"
+
+	"lol.mleku.dev/chk"
 	"lol.mleku.dev/log"
+	"next.orly.dev/pkg/encoders/envelopes"
+	"next.orly.dev/pkg/encoders/envelopes/authenvelope"
+	"next.orly.dev/pkg/encoders/envelopes/closeenvelope"
+	"next.orly.dev/pkg/encoders/envelopes/eventenvelope"
+	"next.orly.dev/pkg/encoders/envelopes/reqenvelope"
 )
 
-func (s *Server) HandleMessage(msg []byte) {
-	log.I.F("received message:\n%s\n", msg)
+func (s *Server) HandleMessage(msg []byte, remote string) {
+	log.D.C(
+		func() string {
+			return fmt.Sprintf(
+				"%s received message:\n%s", remote, msg,
+			)
+		},
+	)
+	var notice []byte
+	var err error
+	var t string
+	var rem []byte
+	if t, rem, err = envelopes.Identify(msg); chk.E(err) {
+		notice = []byte(err.Error())
+	}
+	switch t {
+	case eventenvelope.L:
+		log.D.F("eventenvelope: %s", rem)
+	case reqenvelope.L:
+		log.D.F("reqenvelope: %s", rem)
+	case closeenvelope.L:
+		log.D.F("closeenvelope: %s", rem)
+	case authenvelope.L:
+		log.D.F("authenvelope: %s", rem)
+	default:
+		notice = []byte(fmt.Sprintf("unknown envelope type %s\n%s", t, rem))
+	}
+	if len(notice) > 0 {
+		log.D.C(
+			func() string {
+				return fmt.Sprintf(
+					"notice->%s %s", remote, notice,
+				)
+			},
+		)
+		// if err = noticeenvelope.NewFrom(notice).Write(a.Listener); chk.E(err) {
+		// 	return
+		// }
+	}
+
 }
diff --git a/app/handle-websocket.go b/app/handle-websocket.go
index 84093f9..73d02fa 100644
--- a/app/handle-websocket.go
+++ b/app/handle-websocket.go
@@ -28,6 +28,19 @@ const (
 
 func (s *Server) HandleWebsocket(w http.ResponseWriter, r *http.Request) {
 	remote := GetRemoteFromReq(r)
+	log.D.F("handling websocket connection from %s", remote)
+	if len(s.Config.IPWhitelist) > 0 {
+		for _, ip := range s.Config.IPWhitelist {
+			log.T.F("checking IP whitelist: %s", ip)
+			if strings.HasPrefix(remote, ip) {
+				log.T.F("IP whitelisted %s", remote)
+				goto whitelist
+			}
+		}
+		log.T.F("IP not whitelisted: %s", remote)
+		return
+	}
+whitelist:
 	var cancel context.CancelFunc
 	s.Ctx, cancel = context.WithCancel(s.Ctx)
 	defer cancel()
@@ -39,7 +52,6 @@ func (s *Server) HandleWebsocket(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	defer conn.CloseNow()
-
 	go s.Pinger(s.Ctx, conn, time.NewTicker(time.Second*10), cancel)
 	for {
 		select {
@@ -73,7 +85,7 @@ func (s *Server) HandleWebsocket(w http.ResponseWriter, r *http.Request) {
 			}
 			continue
 		}
-		go s.HandleMessage(msg)
+		go s.HandleMessage(msg, remote)
 	}
 }