update host resolution

Author: totegamma

image.go

@@ -201,7 +201,13 @@ func ImageHandler(c echo.Context) error {
 			return c.String(400, err.Error())
 		}
 
-		targetIPs, err := net.LookupIP(parsedUrl.Host)
+		targetHost := parsedUrl.Host
+		splitHost, _, err := net.SplitHostPort(parsedUrl.Host)
+		if err == nil {
+			targetHost = splitHost
+		}
+
+		targetIPs, err := net.LookupIP(targetHost)
 		if err != nil {
 			err := errors.Wrap(err, "Failed to lookup IP")
 			span.SetAttributes(attribute.String("host", parsedUrl.Host))

summary.go

@@ -76,7 +76,13 @@ func SummaryHandler(c echo.Context) error {
 		return invalidURL(c, "Invalid URL", cacheKey)
 	}
 
-	targetIPs, err := net.LookupIP(parsedUrl.Host)
+	targetHost := parsedUrl.Host
+	splitHost, _, err := net.SplitHostPort(parsedUrl.Host)
+	if err == nil {
+		targetHost = splitHost
+	}
+
+	targetIPs, err := net.LookupIP(targetHost)
 	if err != nil {
 		fmt.Println("Error looking up IP: ", err)
 		return invalidURL(c, parsedUrl.Host, cacheKey)
@@ -92,7 +98,7 @@ func SummaryHandler(c echo.Context) error {
 		for _, targetIP := range targetIPs {
 			if ipnet.Contains(targetIP) {
 				fmt.Println("IP is in deny list: ", targetIP)
-				return invalidURL(c, parsedUrl.Host, cacheKey)
+				return invalidURL(c, targetHost, cacheKey)
 			}
 		}
 	}
@@ -105,7 +111,7 @@ func SummaryHandler(c echo.Context) error {
 	resp, err := client.Do(req)
 	if err != nil {
 		fmt.Println("Error fetching URL: ", err)
-		return invalidURL(c, parsedUrl.Host, cacheKey)
+		return invalidURL(c, targetHost, cacheKey)
 	}
 
 	charset := ""