From adab600e261f0ce6f5bf0a4791d37fe8da49257e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ask=20Bj=C3=B8rn=20Hansen?= <ask@develooper.com>
Date: Sun, 10 Dec 2023 20:55:44 -0800
Subject: [PATCH] Function to get server from IP or ID parameter

---
 server/functions.go   | 46 +++++++++++++++++++++++++++++++++++++++++++
 server/graph_image.go | 26 ++----------------------
 2 files changed, 48 insertions(+), 24 deletions(-)
 create mode 100644 server/functions.go

diff --git a/server/functions.go b/server/functions.go
new file mode 100644
index 0000000..38e75c7
--- /dev/null
+++ b/server/functions.go
@@ -0,0 +1,46 @@
+package server
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"net/netip"
+	"strconv"
+	"time"
+
+	"go.ntppool.org/common/logger"
+	"go.ntppool.org/common/tracing"
+	"go.ntppool.org/data-api/ntpdb"
+)
+
+func (srv *Server) FindServer(ctx context.Context, serverID string) (ntpdb.Server, error) {
+	log := logger.Setup()
+	ctx, span := tracing.Tracer().Start(ctx, "FindServer")
+	defer span.End()
+	q := ntpdb.NewWrappedQuerier(ntpdb.New(srv.db))
+
+	var serverData ntpdb.Server
+	var dberr error
+	if id, err := strconv.Atoi(serverID); id > 0 && err == nil {
+		serverData, dberr = q.GetServerByID(ctx, uint32(id))
+	} else {
+		ip, err := netip.ParseAddr(serverID)
+		if err != nil || !ip.IsValid() {
+			return ntpdb.Server{}, nil // 404 error
+		}
+		serverData, dberr = q.GetServerByIP(ctx, ip.String())
+	}
+	if dberr != nil {
+		if !errors.Is(dberr, sql.ErrNoRows) {
+			log.Error("could not query server id", "err", dberr)
+			return serverData, dberr
+		}
+	}
+
+	if serverData.ID == 0 || (serverData.DeletionOn.Valid && serverData.DeletionOn.Time.Before(time.Now().Add(-1*time.Hour*24*30*24))) {
+		// no data and no error to produce 404 errors
+		return ntpdb.Server{}, nil
+	}
+
+	return serverData, nil
+}
diff --git a/server/graph_image.go b/server/graph_image.go
index 96fbe8b..b0aaca7 100644
--- a/server/graph_image.go
+++ b/server/graph_image.go
@@ -2,16 +2,12 @@ package server
 
 import (
 	"context"
-	"database/sql"
-	"errors"
 	"fmt"
 	"io"
 	"net/http"
 	"net/http/httptrace"
 	"net/url"
 	"os"
-	"strconv"
-	"time"
 
 	"github.com/hashicorp/go-retryablehttp"
 	"github.com/labstack/echo/v4"
@@ -22,8 +18,6 @@ import (
 
 	"go.ntppool.org/common/logger"
 	"go.ntppool.org/common/tracing"
-
-	"go.ntppool.org/data-api/ntpdb"
 )
 
 func (srv *Server) graphImage(c echo.Context) error {
@@ -51,24 +45,8 @@ func (srv *Server) graphImage(c echo.Context) error {
 		return c.Redirect(308, redirectURL.String())
 	}
 
-	q := ntpdb.NewWrappedQuerier(ntpdb.New(srv.db))
-
-	var serverData ntpdb.Server
-	var dberr error
-	if id, err := strconv.Atoi(serverID); id > 0 && err == nil {
-		serverData, dberr = q.GetServerByID(ctx, uint32(id))
-	} else {
-		serverData, dberr = q.GetServerByIP(ctx, serverID)
-	}
-	if dberr != nil {
-		if !errors.Is(dberr, sql.ErrNoRows) {
-			log.Error("could not query server id", "err", dberr)
-			return c.String(http.StatusInternalServerError, "server error")
-		}
-		return c.String(http.StatusNotFound, "not found")
-	}
-
-	if serverData.ID == 0 || (serverData.DeletionOn.Valid && serverData.DeletionOn.Time.Before(time.Now())) {
+	serverData, err := srv.FindServer(ctx, serverID)
+	if serverData.ID == 0 {
 		return c.String(http.StatusNotFound, "not found")
 	}