From fab091d75a35a764a4e0d3dd15799ba1fe33b3a6 Mon Sep 17 00:00:00 2001
From: memory <github@memory.blank.org>
Date: Fri, 18 May 2018 20:36:14 -0400
Subject: [PATCH] add healthcheck handling (#106)

Add a special path -- `/ok` that can be used as a healthcheck for e.g.
kubernetes or amazon ECS.

When not in debug mode, do not generate logs for the healthcheck path.

This requires implementing our own version of ginrus.Ginrus, as the
upstream one does not support the `notlogged` argument(s) that
gin.LoggerWithWriter has.
---
 handlers/handlers.go | 66 ++++++++++++++++++++++++++++++++++++++++++--
 handlers/public.go   | 10 +++++++
 2 files changed, 74 insertions(+), 2 deletions(-)

diff --git a/handlers/handlers.go b/handlers/handlers.go
index 154155c..6d46eae 100644
--- a/handlers/handlers.go
+++ b/handlers/handlers.go
@@ -6,7 +6,6 @@ import (
 	"net/http"
 	"time"
 
-	"github.com/gin-gonic/contrib/ginrus"
 	"github.com/sirupsen/logrus"
 
 	"github.com/gin-gonic/gin"
@@ -27,6 +26,63 @@ type Handler struct {
 // DoNotPrivateKeyChecking is used for testing
 var DoNotPrivateKeyChecking = false
 
+type loggerEntryWithFields interface {
+	WithFields(fields logrus.Fields) *logrus.Entry
+}
+
+// Ginrus returns a gin.HandlerFunc (middleware) that logs requests using logrus.
+//
+// Requests with errors are logged using logrus.Error().
+// Requests without errors are logged using logrus.Info().
+//
+// It receives:
+//   1. A time package format string (e.g. time.RFC3339).
+//   2. A boolean stating whether to use UTC time zone or local.
+//   3. Optionally, a list of paths to skip logging for (this is why
+//      we are not using upstream github.com/gin-gonic/contrib/ginrus)
+func Ginrus(logger loggerEntryWithFields, timeFormat string, utc bool, notlogged ...string) gin.HandlerFunc {
+	var skip map[string]struct{}
+	if length := len(notlogged); length > 0 {
+		skip = make(map[string]struct{}, length)
+		for _, path := range notlogged {
+			skip[path] = struct{}{}
+		}
+	}
+
+	return func(c *gin.Context) {
+		start := time.Now()
+		// some evil middlewares modify this values
+		path := c.Request.URL.Path
+		c.Next()
+
+		// log only when path is not being skipped
+		if _, ok := skip[path]; !ok {
+			end := time.Now()
+			latency := end.Sub(start)
+			if utc {
+				end = end.UTC()
+			}
+
+			entry := logger.WithFields(logrus.Fields{
+				"status":     c.Writer.Status(),
+				"method":     c.Request.Method,
+				"path":       path,
+				"ip":         c.ClientIP(),
+				"latency":    latency,
+				"user-agent": c.Request.UserAgent(),
+				"time":       end.Format(timeFormat),
+			})
+
+			if len(c.Errors) > 0 {
+				// Append error field if this is an erroneous request.
+				entry.Error(c.Errors.String())
+			} else {
+				entry.Info()
+			}
+		}
+	}
+}
+
 // New initializes the http handlers
 func New(store stores.Store) (*Handler, error) {
 	if !util.GetConfig().EnableDebugMode {
@@ -78,7 +134,12 @@ func (h *Handler) setHandlers() error {
 	if err := h.addTemplatesFromFS([]string{"token.html", "protected.html"}); err != nil {
 		return errors.Wrap(err, "could not add templates from FS")
 	}
-	h.engine.Use(ginrus.Ginrus(logrus.StandardLogger(), time.RFC3339, false))
+	if !util.GetConfig().EnableDebugMode {
+		// if we are not in debug mode, do not log healthchecks
+		h.engine.Use(Ginrus(logrus.StandardLogger(), time.RFC3339, false, "/ok"))
+	} else {
+		h.engine.Use(Ginrus(logrus.StandardLogger(), time.RFC3339, false))
+	}
 	protected := h.engine.Group("/api/v1/protected")
 	if util.GetConfig().AuthBackend == "oauth" {
 		logrus.Info("Using OAuth auth backend")
@@ -96,6 +157,7 @@ func (h *Handler) setHandlers() error {
 
 	h.engine.GET("/api/v1/info", h.handleInfo)
 	h.engine.GET("/d/:id/:hash", h.handleDelete)
+	h.engine.GET("/ok", h.handleHealthcheck)
 
 	// Handling the shorted URLs, if no one exists, it checks
 	// in the filesystem and sets headers for caching
diff --git a/handlers/public.go b/handlers/public.go
index d3a495c..f1acd44 100644
--- a/handlers/public.go
+++ b/handlers/public.go
@@ -139,6 +139,16 @@ func (h *Handler) handleGetVisitors(c *gin.Context) {
 	c.JSON(http.StatusOK, dataSets)
 }
 
+// handleHealthcheck returns success for healthcheckers without polluting logs
+func (h *Handler) handleHealthcheck(c *gin.Context) {
+	out := struct {
+		Status string `json:"status"`
+	}{
+		"OK",
+	}
+	c.JSON(http.StatusOK, out)
+}
+
 func (h *Handler) handleInfo(c *gin.Context) {
 	out := struct {
 		util.Info