From 2ee24c8d510eb8453ab9814ea216a1b1ff63b95d Mon Sep 17 00:00:00 2001
From: bagas <bagas@fossy.my.id>
Date: Sun, 25 Jan 2026 18:17:37 +0700
Subject: [PATCH] test(config): add test for keyloc and header size

---
 internal/config/config.go      |  2 ++
 internal/config/config_test.go | 30 ++++++++++++++++++++++++++++++
 internal/config/loader.go      | 15 ++++++++++++++-
 3 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/internal/config/config.go b/internal/config/config.go
index 19bbc49..3e6c9e1 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -23,6 +23,7 @@ type Config interface {
 	AllowedPortsEnd() uint16
 
 	BufferSize() int
+	HeaderSize() int
 
 	PprofEnabled() bool
 	PprofPort() string
@@ -60,6 +61,7 @@ func (c *config) ACMEStaging() bool         { return c.acmeStaging }
 func (c *config) AllowedPortsStart() uint16 { return c.allowedPortsStart }
 func (c *config) AllowedPortsEnd() uint16   { return c.allowedPortsEnd }
 func (c *config) BufferSize() int           { return c.bufferSize }
+func (c *config) HeaderSize() int           { return c.headerSize }
 func (c *config) PprofEnabled() bool        { return c.pprofEnabled }
 func (c *config) PprofPort() string         { return c.pprofPort }
 func (c *config) Mode() types.ServerMode    { return c.mode }
diff --git a/internal/config/config_test.go b/internal/config/config_test.go
index effbcaf..b102094 100644
--- a/internal/config/config_test.go
+++ b/internal/config/config_test.go
@@ -188,6 +188,32 @@ func TestParseBufferSize(t *testing.T) {
 	}
 }
 
+func TestParseHeaderSize(t *testing.T) {
+	tests := []struct {
+		name   string
+		val    string
+		expect int
+	}{
+		{"valid size", "8192", 8192},
+		{"default size", "", 4096},
+		{"too small", "1024", 4096},
+		{"too large", "2000000", 4096},
+		{"invalid format", "abc", 4096},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.val != "" {
+				t.Setenv("MAX_HEADER_SIZE", tt.val)
+			} else {
+				os.Unsetenv("MAX_HEADER_SIZE")
+			}
+			size := parseHeaderSize()
+			assert.Equal(t, tt.expect, size)
+		})
+	}
+}
+
 func TestParse(t *testing.T) {
 	tests := []struct {
 		name      string
@@ -271,6 +297,7 @@ func TestGetters(t *testing.T) {
 		"PORT":             "2222",
 		"HTTP_PORT":        "80",
 		"HTTPS_PORT":       "443",
+		"KEY_LOC":          "certs/ssh/id_rsa",
 		"TLS_ENABLED":      "true",
 		"TLS_REDIRECT":     "true",
 		"TLS_STORAGE_PATH": "certs/tls/",
@@ -279,6 +306,7 @@ func TestGetters(t *testing.T) {
 		"ACME_STAGING":     "true",
 		"ALLOWED_PORTS":    "1000-2000",
 		"BUFFER_SIZE":      "16384",
+		"MAX_HEADER_SIZE":  "4096",
 		"PPROF_ENABLED":    "true",
 		"PPROF_PORT":       "7070",
 		"MODE":             "standalone",
@@ -299,6 +327,7 @@ func TestGetters(t *testing.T) {
 	assert.Equal(t, "2222", cfg.SSHPort())
 	assert.Equal(t, "80", cfg.HTTPPort())
 	assert.Equal(t, "443", cfg.HTTPSPort())
+	assert.Equal(t, "certs/ssh/id_rsa", cfg.KeyLoc())
 	assert.Equal(t, true, cfg.TLSEnabled())
 	assert.Equal(t, true, cfg.TLSRedirect())
 	assert.Equal(t, "certs/tls/", cfg.TLSStoragePath())
@@ -308,6 +337,7 @@ func TestGetters(t *testing.T) {
 	assert.Equal(t, uint16(1000), cfg.AllowedPortsStart())
 	assert.Equal(t, uint16(2000), cfg.AllowedPortsEnd())
 	assert.Equal(t, 16384, cfg.BufferSize())
+	assert.Equal(t, 4096, cfg.HeaderSize())
 	assert.Equal(t, true, cfg.PprofEnabled())
 	assert.Equal(t, "7070", cfg.PprofPort())
 	assert.Equal(t, types.ServerMode(types.ServerModeSTANDALONE), cfg.Mode())
diff --git a/internal/config/loader.go b/internal/config/loader.go
index a1fde22..5cbfe1f 100644
--- a/internal/config/loader.go
+++ b/internal/config/loader.go
@@ -31,7 +31,8 @@ type config struct {
 	allowedPortsEnd   uint16
 
 	bufferSize int
-
+	headerSize int
+	
 	pprofEnabled bool
 	pprofPort    string
 
@@ -73,6 +74,7 @@ func parse() (*config, error) {
 	}
 
 	bufferSize := parseBufferSize()
+	headerSize := parseHeaderSize()
 
 	pprofEnabled := getenvBool("PPROF_ENABLED", false)
 	pprofPort := getenv("PPROF_PORT", "6060")
@@ -100,6 +102,7 @@ func parse() (*config, error) {
 		allowedPortsStart: start,
 		allowedPortsEnd:   end,
 		bufferSize:        bufferSize,
+		headerSize:        headerSize,
 		pprofEnabled:      pprofEnabled,
 		pprofPort:         pprofPort,
 		mode:              mode,
@@ -161,6 +164,16 @@ func parseBufferSize() int {
 	return size
 }
 
+func parseHeaderSize() int {
+	raw := getenv("MAX_HEADER_SIZE", "4096")
+	size, err := strconv.Atoi(raw)
+	if err != nil || size < 4096 || size > 131072 {
+		log.Println("Invalid BUFFER_SIZE, falling back to 4096")
+		return 4096
+	}
+	return size
+}
+
 func getenv(key, def string) string {
 	if v := os.Getenv(key); v != "" {
 		return v