Add per site ttl settings

deflect-ca · Jun 3, 2024 · b237808 · b237808
1 parent 86086a0
commit b237808
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 2 deletions.
diff --git a/internal/config.go b/internal/config.go
@@ -49,6 +49,8 @@ type Config struct {
 	ExpiringDecisionTtlSeconds             int                 `yaml:"expiring_decision_ttl_seconds"`
 	BlockIPTtlSeconds                      int                 `yaml:"block_ip_ttl_seconds"`
 	BlockSessionTtlSeconds                 int                 `yaml:"block_session_ttl_seconds"`
+	SitesToBlockIPTtlSeconds               map[string]int      `yaml:"sites_to_block_ip_ttl_seconds"`
+	SitesToBlockSessionTtlSeconds          map[string]int      `yaml:"sites_to_block_session_ttl_seconds"`
 	TooManyFailedChallengesIntervalSeconds int                 `yaml:"too_many_failed_challenges_interval_seconds"`
 	TooManyFailedChallengesThreshold       int                 `yaml:"too_many_failed_challenges_threshold"`
 	PasswordCookieTtlSeconds               int                 `yaml:"password_cookie_ttl_seconds"`

diff --git a/internal/kafka.go b/internal/kafka.go
@@ -137,6 +137,24 @@ func RunKafkaReader(
 	}
 }
 
+func getBlockIpTtl(config *Config, host string) (blockIpTtl int) {
+	blockIpTtl = config.BlockSessionTtlSeconds
+	if ttl, ok := config.SitesToBlockIPTtlSeconds[host]; ok {
+		log.Printf("KAFKA: found site-specific block_ip ttl %s %d\n", host, ttl)
+		blockIpTtl = ttl
+	}
+	return
+}
+
+func getBlockSessionTtl(config *Config, host string) (blockSessionTtl int) {
+	blockSessionTtl = config.BlockIPTtlSeconds
+	if ttl, ok := config.SitesToBlockSessionTtlSeconds[host]; ok {
+		log.Printf("KAFKA: found site-specific block_session ttl %s %d\n", host, ttl)
+		blockSessionTtl = ttl
+	}
+	return
+}
+
 func handleCommand(
 	config *Config,
 	command commandMessage,
@@ -155,13 +173,15 @@ func handleCommand(
 		handleIPCommand(config, command, decisionListsMutex, decisionLists, Challenge, config.ExpiringDecisionTtlSeconds)
 		break
 	case "block_ip":
-		handleIPCommand(config, command, decisionListsMutex, decisionLists, NginxBlock, config.BlockIPTtlSeconds)
+		ttl := getBlockIpTtl(config, command.Host)
+		handleIPCommand(config, command, decisionListsMutex, decisionLists, NginxBlock, ttl)
 		break
 	case "challenge_session":
 		handleSessionCommand(config, command, decisionListsMutex, decisionLists, Challenge, config.ExpiringDecisionTtlSeconds)
 		break
 	case "block_session":
-		handleSessionCommand(config, command, decisionListsMutex, decisionLists, NginxBlock, config.BlockSessionTtlSeconds)
+		ttl := getBlockSessionTtl(config, command.Host)
+		handleSessionCommand(config, command, decisionListsMutex, decisionLists, NginxBlock, ttl)
 		break
 	default:
 		log.Printf("KAFKA: unrecognized command name: %s\n", command.Name)