From 65d69fd692889839f36e9171c7ede4477a8a4ce2 Mon Sep 17 00:00:00 2001
From: lib0xidium <oxidium@oxide.rs>
Date: Thu, 10 Oct 2024 11:49:09 +0200
Subject: [PATCH] fix: getHostByName

- changed getHostByName
- added unit-test
---
 network.go      | 16 +++++++++++-----
 network_test.go | 33 ++++++++++++++++++++++++++++-----
 2 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/network.go b/network.go
index 108d78a9..5442ca21 100644
--- a/network.go
+++ b/network.go
@@ -1,12 +1,18 @@
 package sprig
 
 import (
-	"math/rand"
 	"net"
 )
 
-func getHostByName(name string) string {
-	addrs, _ := net.LookupHost(name)
-	//TODO: add error handing when release v3 comes out
-	return addrs[rand.Intn(len(addrs))]
+const NUM_TRIES = 3
+
+func getHostByName(name string) ([]string, error) {
+	err := error(nil)
+	addrs := []string(nil)
+	for tries := 0; tries < NUM_TRIES; tries++ {
+		if addrs, err = net.LookupHost(name); err == nil {
+			return addrs, nil
+		}
+	}
+	return addrs, err
 }
diff --git a/network_test.go b/network_test.go
index 9c153f0a..0b222be5 100644
--- a/network_test.go
+++ b/network_test.go
@@ -2,17 +2,40 @@ package sprig
 
 import (
 	"net"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 )
 
 func TestGetHostByName(t *testing.T) {
-	tpl := `{{"www.google.com" | getHostByName}}`
+	// GIVEN a valid hostname
+	tpl := `{{"google.com" | getHostByName}}`
 
-	resolvedIP, _ := runRaw(tpl, nil)
+	// WHEN getHostByName is executed
+	resolvedIP, err := runRaw(tpl, nil)
 
-	ip := net.ParseIP(resolvedIP)
-	assert.NotNil(t, ip)
-	assert.NotEmpty(t, ip)
+	// THEN the resolved IP should not be empty and no error should be returned
+	assert.NotEmpty(t, resolvedIP)
+	assert.NoError(t, err)
+
+	// result has type string, but it should be a slice of strings
+	// convert it to a slice of strings
+	resolvedIPs := strings.Split(resolvedIP[1:len(resolvedIP)-1], " ")
+
+	// Check if the resolved IP is a valid IP address
+	parsedIP := net.ParseIP(resolvedIPs[0])
+	assert.NotNil(t, parsedIP)
+}
+
+func TestGetHostByNameNXDomain(t *testing.T) {
+	// GIVEN an invalid hostname
+	tpl := `{{"invalid.invalid" | getHostByName}}`
+
+	// WHEN getHostByName is executed
+	resolvedIP, err := runRaw(tpl, nil)
+
+	// THEN the resolved IP should be empty and an error should be returned
+	assert.Empty(t, resolvedIP)
+	assert.Error(t, err)
 }