heicode-mananger/heicode/middleware/device_auth_diagnostic_test.go

package middleware

import (
	"errors"
	"net/http/httptest"
	"strconv"
	"testing"
	"time"

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

// V2 auth diagnostic helpers — covers the client-facing contract that
// every V2 401 carries both X-Heicode-Server-Time (so the client can
// detect clock drift) and X-Heicode-Auth-Error (so it can map the failure
// to a localized message instead of "Token invalid").
//
// These tests pin the public error-code strings; renaming any of them
// without coordinated client-side rollout will break drift detection on
// already-shipped desktop builds.

func TestV2AuthErrorCode_KnownSentinels(t *testing.T) {
	cases := []struct {
		err  error
		want string
	}{
		{ErrDeviceTimestampOutOfWindow, V2AuthErrorTimestampDrift},
		{ErrDeviceNonceReplayed, V2AuthErrorNonceReplay},
		{ErrDeviceRevoked, V2AuthErrorRevoked},
		{ErrDeviceFingerprintMismatch, V2AuthErrorFingerprintMismatch},
		{ErrDeviceSignatureInvalid, V2AuthErrorSignatureInvalid},
		{ErrDeviceSignatureMissing, V2AuthErrorSignatureMissing},
		{ErrV2DeviceNotFound, V2AuthErrorDeviceNotFound},
		{ErrV2EphPubkeyMissing, V2AuthErrorEphPubkeyMissing},
		{ErrV2EphPubkeyMalformed, V2AuthErrorEphPubkeyMalformed},
		{ErrV2AEADDecryptFailed, V2AuthErrorBodyDecryptFailed},
		{ErrV2BodyTooShort, V2AuthErrorBodyTooShort},
		{ErrV2BodyReadFailed, V2AuthErrorBodyReadFailed},
		{ErrDeviceSignatureBodyReadFail, V2AuthErrorBodyReadFailed},
		{ErrV2KeyDeriveFailed, V2AuthErrorKeyDeriveFailed},
		{ErrV2HelpersNotInitialized, V2AuthErrorServerNotReady},
	}
	for _, tc := range cases {
		got := V2AuthErrorCode(tc.err)
		if got != tc.want {
			t.Errorf("V2AuthErrorCode(%v) = %q, want %q", tc.err, got, tc.want)
		}
	}
}

func TestV2AuthErrorCode_UnknownDefaultsToUnknown(t *testing.T) {
	// Any error not in the sentinel list (e.g. a wrapped DB error,
	// gorm.ErrRecordNotFound, network read) must map to "unknown"
	// rather than panicking or leaking the Go error string to the
	// client. The client treats "unknown" as a generic auth failure.
	custom := errors.New("not a v2 sentinel")
	if got := V2AuthErrorCode(custom); got != V2AuthErrorUnknown {
		t.Fatalf("expected %q for unknown error, got %q", V2AuthErrorUnknown, got)
	}
}

func TestSetV2AuthDiagnosticHeaders_SetsBothHeaders(t *testing.T) {
	gin.SetMode(gin.TestMode)
	rec := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(rec)
	c.Request = httptest.NewRequest("POST", "/v1/messages", nil)

	before := time.Now().UnixMilli()
	SetV2AuthDiagnosticHeaders(c, ErrDeviceTimestampOutOfWindow)
	after := time.Now().UnixMilli()

	got := rec.Header().Get(HeaderAuthError)
	if got != V2AuthErrorTimestampDrift {
		t.Errorf("X-Heicode-Auth-Error = %q, want %q", got, V2AuthErrorTimestampDrift)
	}

	tsStr := rec.Header().Get(HeaderServerTime)
	if tsStr == "" {
		t.Fatalf("X-Heicode-Server-Time missing")
	}
	ts, err := strconv.ParseInt(tsStr, 10, 64)
	if err != nil {
		t.Fatalf("X-Heicode-Server-Time %q not unix_ms: %v", tsStr, err)
	}
	if ts < before || ts > after {
		t.Errorf("X-Heicode-Server-Time = %d, want in [%d, %d]", ts, before, after)
	}
}

func TestSetV2AuthDiagnosticHeaders_UnknownErrorStillStampsTime(t *testing.T) {
	// Even when the error is unknown, the server MUST still set the
	// time header — the client uses it to detect "my clock is way off"
	// even if it can't act on the error code itself.
	gin.SetMode(gin.TestMode)
	rec := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(rec)
	c.Request = httptest.NewRequest("POST", "/v1/messages", nil)

	SetV2AuthDiagnosticHeaders(c, errors.New("opaque"))

	if got := rec.Header().Get(HeaderAuthError); got != V2AuthErrorUnknown {
		t.Errorf("want %q, got %q", V2AuthErrorUnknown, got)
	}
	if rec.Header().Get(HeaderServerTime) == "" {
		t.Fatalf("X-Heicode-Server-Time missing on unknown-error path")
	}
}