warmbly/internal/api/middleware/node_heartbeat_auth_test.go

package middleware

import (
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"

	"github.com/gin-gonic/gin"
)

// heartbeatOutcome sends one heartbeat with bearer and reports the status and whether it was update-only.
func heartbeatOutcome(t *testing.T, internal, broker, accept, bearer string) (int, bool) {
	t.Helper()
	t.Setenv("INTERNAL_API_TOKEN", internal)
	t.Setenv("NODE_BROKER_TOKEN", broker)
	t.Setenv("NODE_ACCEPT_INTERNAL_TOKEN", accept)
	brokerTokenOnce = sync.Once{}
	brokerToken = nil
	updateOnly := false
	h := &Handler{}
	r := gin.New()
	r.POST("/fleet/heartbeat", h.NodeHeartbeatAuthMiddleware(), func(c *gin.Context) {
		updateOnly = IsNodeUpdateOnly(c)
		c.Status(http.StatusOK)
	})
	req := httptest.NewRequest(http.MethodPost, "/fleet/heartbeat", nil)
	if bearer != "" {
		req.Header.Set("Authorization", "Bearer "+bearer)
	}
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w.Code, updateOnly
}

func TestNodeHeartbeatAuth(t *testing.T) {
	for _, tc := range []struct {
		name                            string
		internal, broker, accept, token string
		code                            int
		updateOnly                      bool
	}{
		{"broker token is a full beat", "shared", "broker", "", "broker", http.StatusOK, false},
		{"shared token is update-only", "shared", "broker", "", "shared", http.StatusOK, true},
		{"shared token is a full beat while accepted", "shared", "broker", "true", "shared", http.StatusOK, false},
		{"single-token instance is a full beat", "shared", "", "", "shared", http.StatusOK, false},
		{"same value for both is a full beat", "same", "same", "", "same", http.StatusOK, false},
		{"wrong token is refused", "shared", "broker", "", "other", http.StatusUnauthorized, false},
		{"no token is refused", "shared", "broker", "", "", http.StatusUnauthorized, false},
		{"nothing configured is refused", "", "", "", "anything", http.StatusUnauthorized, false},
	} {
		t.Run(tc.name, func(t *testing.T) {
			code, updateOnly := heartbeatOutcome(t, tc.internal, tc.broker, tc.accept, tc.token)
			if code != tc.code || updateOnly != tc.updateOnly {
				t.Fatalf("got %d update-only=%v, want %d update-only=%v", code, updateOnly, tc.code, tc.updateOnly)
			}
		})
	}
}