From 672547cd07d0620ae988ec2ff96ac298bcb07778 Mon Sep 17 00:00:00 2001
From: Svilen Markov <7613769+svilenmarkov@users.noreply.github.com>
Date: Mon, 30 Sep 2024 00:24:44 +0100
Subject: [PATCH] Switch to using custom type for pihole's top blocked domains

---
 internal/feed/pihole.go | 62 +++++++++++++++++++----------------------
 1 file changed, 29 insertions(+), 33 deletions(-)

diff --git a/internal/feed/pihole.go b/internal/feed/pihole.go
index 5a4b7f8..3c7f1b5 100644
--- a/internal/feed/pihole.go
+++ b/internal/feed/pihole.go
@@ -1,6 +1,7 @@
 package feed
 
 import (
+	"encoding/json"
 	"errors"
 	"log/slog"
 	"net/http"
@@ -9,24 +10,33 @@ import (
 )
 
 type piholeStatsResponse struct {
-	TotalQueries      int            `json:"dns_queries_today"`
-	QueriesSeries     map[int64]int  `json:"domains_over_time"`
-	BlockedQueries    int            `json:"ads_blocked_today"`
-	BlockedSeries     map[int64]int  `json:"ads_over_time"`
-	BlockedPercentage float64        `json:"ads_percentage_today"`
-	TopBlockedDomains map[string]int `json:"top_ads"`
-	DomainsBlocked    int            `json:"domains_being_blocked"`
+	TotalQueries      int                     `json:"dns_queries_today"`
+	QueriesSeries     map[int64]int           `json:"domains_over_time"`
+	BlockedQueries    int                     `json:"ads_blocked_today"`
+	BlockedSeries     map[int64]int           `json:"ads_over_time"`
+	BlockedPercentage float64                 `json:"ads_percentage_today"`
+	TopBlockedDomains piholeTopBlockedDomains `json:"top_ads"`
+	DomainsBlocked    int                     `json:"domains_being_blocked"`
 }
 
 // If user has some level of privacy enabled on Pihole, `json:"top_ads"` is an empty array
-// Use alternate struct without that field to avoid error when unmarshalling
-type piholeStatsResponsePrivate struct {
-	TotalQueries      int            `json:"dns_queries_today"`
-	QueriesSeries     map[int64]int  `json:"domains_over_time"`
-	BlockedQueries    int            `json:"ads_blocked_today"`
-	BlockedSeries     map[int64]int  `json:"ads_over_time"`
-	BlockedPercentage float64        `json:"ads_percentage_today"`
-	DomainsBlocked    int            `json:"domains_being_blocked"`
+// Use custom unmarshal behavior to avoid not getting the rest of the valid data when unmarshalling
+type piholeTopBlockedDomains map[string]int
+
+func (p *piholeTopBlockedDomains) UnmarshalJSON(data []byte) error {
+	// NOTE: do not change to piholeTopBlockedDomains type here or it will cause a stack overflow
+	// because of the UnmarshalJSON method getting called recursively
+	temp := make(map[string]int)
+
+	err := json.Unmarshal(data, &temp)
+
+	if err != nil {
+		*p = make(piholeTopBlockedDomains)
+	} else {
+		*p = temp
+	}
+
+	return nil
 }
 
 func FetchPiholeStats(instanceURL, token string) (*DNSStats, error) {
@@ -42,27 +52,13 @@ func FetchPiholeStats(instanceURL, token string) (*DNSStats, error) {
 	if err != nil {
 		return nil, err
 	}
-	
+
 	responseJson, err := decodeJsonFromRequest[piholeStatsResponse](defaultClient, request)
-	
+
 	if err != nil {
-		// Refer to piholeStatsResponsePrivate above
-		responseJsonPriv, err :=
-			decodeJsonFromRequest[piholeStatsResponsePrivate](defaultClient, request)
-		if err != nil {
-			return nil, err
-		}
-		
-		// Copy the results back to responseJson, leaving the TopBlockedDomains field empty
-		responseJson.TotalQueries = responseJsonPriv.TotalQueries
-		responseJson.QueriesSeries = responseJsonPriv.QueriesSeries
-		responseJson.BlockedQueries = responseJsonPriv.BlockedQueries
-		responseJson.BlockedSeries = responseJsonPriv.BlockedSeries
-		responseJson.BlockedPercentage = responseJsonPriv.BlockedPercentage
-		responseJson.TopBlockedDomains = make(map[string]int)
-		responseJson.DomainsBlocked = responseJsonPriv.DomainsBlocked
+		return nil, err
 	}
-	
+
 	stats := &DNSStats{
 		TotalQueries:   responseJson.TotalQueries,
 		BlockedQueries: responseJson.BlockedQueries,