fmt (#1640)

Author: hendoxc

pkg/beholder/auth.go

@@ -129,24 +129,16 @@ func (r *rotatingAuth) Headers(ctx context.Context) (map[string]string, error) {
 			return returnHeader, nil
 		}
 
-		// Append the bytes of the public key with bytes of the timestamp to create the message to sign
-		ts := time.Now()
-		tsBytes := make([]byte, 8)
-		binary.BigEndian.PutUint64(tsBytes, uint64(ts.UnixNano()))
-		msgBytes := append(r.csaPubKey, tsBytes...)
-
 		ctxWithTimeout, cancel := context.WithTimeout(ctx, r.signerTimeout)
 		defer cancel()
 
-		// Sign(public key bytes + timestamp bytes)
-		signature, err := r.signer.Sign(ctxWithTimeout, fmt.Sprintf("%x", r.csaPubKey), msgBytes)
+		ts := time.Now()
+
+		newHeaders, err := NewAuthHeaderV2(ctxWithTimeout, r.csaPubKey, r.signer, ts)
 		if err != nil {
-			return nil, fmt.Errorf("beholder: failed to sign auth header: %w", err)
+			return nil, fmt.Errorf("beholder: failed to create auth header: %w", err)
 		}
 
-		newHeaders := make(map[string]string)
-		newHeaders[authHeaderKey] = fmt.Sprintf("%s:%x:%d:%x", authHeaderV2, r.csaPubKey, ts.UnixNano(), signature)
-
 		r.headers.Store(newHeaders)
 		r.lastUpdatedNanos.Store(ts.UnixNano())
 	}
@@ -200,6 +192,32 @@ func NewAuthHeaders(ed25519Signer crypto.Signer) (map[string]string, error) {
 	return map[string]string{authHeaderKey: headerValue}, nil
 }
 
+// NewAuthHeadersV2 creates the V2 format of the auth header value to be included on requests.
+// This format includes a timestamp as part of the message to sign.
+// The current format for V2 headers is:
+//
+// <version>:<public_key_hex>:<timestamp_bytes>:<signature_hex>
+//
+// where the byte value of <public_key_hex> + <timestamp_bytes> is what's being signed
+func NewAuthHeaderV2(ctx context.Context, pubKey ed25519.PublicKey, signer Signer, ts time.Time) (map[string]string, error) {
+
+	// Append the bytes of the public key with bytes of the timestamp to create the message to sign
+	tsBytes := make([]byte, 8)
+	binary.BigEndian.PutUint64(tsBytes, uint64(ts.UnixNano()))
+	msgBytes := append(pubKey, tsBytes...)
+
+	// Sign(public key bytes + timestamp bytes)
+	signature, err := signer.Sign(ctx, fmt.Sprintf("%x", pubKey), msgBytes)
+	if err != nil {
+		return nil, fmt.Errorf("beholder: failed to sign auth header: %w", err)
+	}
+
+	headers := make(map[string]string)
+	headers[authHeaderKey] = fmt.Sprintf("%s:%x:%d:%x", authHeaderV2, pubKey, ts.UnixNano(), signature)
+
+	return headers, nil
+}
+
 func authDialOpt(auth PerRPCCredentialsProvider) grpc.DialOption {
 	return grpc.WithPerRPCCredentials(auth.Credentials())
 }

pkg/beholder/auth_test.go

@@ -3,6 +3,7 @@ package beholder_test
 import (
 	"context"
 	"crypto/ed25519"
+	"encoding/binary"
 	"encoding/hex"
 	"fmt"
 	"strings"
@@ -34,6 +35,165 @@ func TestBuildAuthHeaders(t *testing.T) {
 	assert.Equal(t, expectedHeaders, headers)
 }
 
+func TestNewAuthHeaderV2(t *testing.T) {
+	// Generate test key pair
+	pubKey, privKey, err := ed25519.GenerateKey(nil)
+	require.NoError(t, err)
+
+	t.Run("creates valid V2 auth headers", func(t *testing.T) {
+		mockSigner := &MockSigner{}
+
+		ts := time.Now()
+
+		// Create the expected message bytes (pubkey + timestamp)
+		expectedSignature := []byte("test-signature")
+		mockSigner.
+			On("Sign", t.Context(), hex.EncodeToString(pubKey), mock.Anything).
+			Return(expectedSignature, nil).
+			Once()
+
+		headers, err := beholder.NewAuthHeaderV2(t.Context(), pubKey, mockSigner, ts)
+		require.NoError(t, err)
+		require.NotNil(t, headers)
+		require.Contains(t, headers, "X-Beholder-Node-Auth-Token")
+
+		authHeader := headers["X-Beholder-Node-Auth-Token"]
+		parts := strings.Split(authHeader, ":")
+		require.Len(t, parts, 4, "Auth header should have format version:pubkey_hex:timestamp:signature_hex")
+
+		assert.Equal(t, "2", parts[0], "Version should be 2")
+		assert.Equal(t, hex.EncodeToString(pubKey), parts[1], "Public key should match")
+		assert.Equal(t, fmt.Sprintf("%d", ts.UnixNano()), parts[2], "Timestamp should match")
+		assert.Equal(t, hex.EncodeToString(expectedSignature), parts[3], "Signature should match")
+
+		mockSigner.AssertExpectations(t)
+	})
+	t.Run("returns error when signer fails", func(t *testing.T) {
+		mockSigner := &MockSigner{}
+		ts := time.Now()
+
+		expectedErr := fmt.Errorf("signing failed")
+		mockSigner.
+			On("Sign", t.Context(), hex.EncodeToString(pubKey), mock.Anything).
+			Return([]byte{}, expectedErr).
+			Once()
+
+		headers, err := beholder.NewAuthHeaderV2(t.Context(), pubKey, mockSigner, ts)
+		require.Error(t, err)
+		assert.Nil(t, headers)
+		assert.Contains(t, err.Error(), "beholder: failed to sign auth header")
+		assert.Contains(t, err.Error(), expectedErr.Error())
+
+		mockSigner.AssertExpectations(t)
+	})
+
+	t.Run("verifies signature with ed25519", func(t *testing.T) {
+		// Use a real signature for verification
+		mockSigner := &MockSigner{}
+		ts := time.Now()
+
+		// Calculate the message that should be signed
+		tsBytes := make([]byte, 8)
+		binary.BigEndian.PutUint64(tsBytes, uint64(ts.UnixNano()))
+		msgBytes := append(pubKey, tsBytes...)
+
+		// Sign with the actual private key
+		realSignature := ed25519.Sign(privKey, msgBytes)
+
+		mockSigner.
+			On("Sign", t.Context(), hex.EncodeToString(pubKey), mock.MatchedBy(func(data []byte) bool {
+				// Match if the data contains pubkey + timestamp
+				return len(data) == len(pubKey)+8 && string(data[:len(pubKey)]) == string(pubKey)
+			})).
+			Return(realSignature, nil).
+			Once()
+
+		headers, err := beholder.NewAuthHeaderV2(t.Context(), pubKey, mockSigner, ts)
+		require.NoError(t, err)
+		require.NotNil(t, headers)
+
+		authHeader := headers["X-Beholder-Node-Auth-Token"]
+		parts := strings.Split(authHeader, ":")
+		require.Len(t, parts, 4)
+
+		signatureBytes, err := hex.DecodeString(parts[3])
+		require.NoError(t, err)
+
+		// Verify the signature
+		valid := ed25519.Verify(pubKey, msgBytes, signatureBytes)
+		assert.True(t, valid, "Signature should be valid")
+
+		mockSigner.AssertExpectations(t)
+	})
+
+	t.Run("handles context cancellation", func(t *testing.T) {
+		mockSigner := &MockSigner{}
+
+		ts := time.Now()
+
+		mockSigner.
+			On("Sign", t.Context(), hex.EncodeToString(pubKey), mock.Anything).
+			Return([]byte{}, context.Canceled).
+			Maybe()
+
+		headers, err := beholder.NewAuthHeaderV2(t.Context(), pubKey, mockSigner, ts)
+
+		// The function should propagate the context error
+		if err != nil {
+			assert.Contains(t, err.Error(), "beholder: failed to sign auth header")
+		}
+
+		// If mockSigner.Sign was called and returned error, headers should be nil
+		if err != nil {
+			assert.Nil(t, headers)
+		}
+	})
+
+	t.Run("uses correct keyID format", func(t *testing.T) {
+		mockSigner := &MockSigner{}
+		ts := time.Now()
+
+		var capturedKeyID string
+		mockSigner.
+			On("Sign", t.Context(), mock.Anything, mock.Anything).
+			Run(func(args mock.Arguments) {
+				capturedKeyID = args.Get(1).(string)
+			}).
+			Return([]byte("signature"), nil).
+			Once()
+
+		_, err := beholder.NewAuthHeaderV2(t.Context(), pubKey, mockSigner, ts)
+		require.NoError(t, err)
+
+		// Verify keyID is hex-encoded public key
+		assert.Equal(t, hex.EncodeToString(pubKey), capturedKeyID)
+
+		mockSigner.AssertExpectations(t)
+	})
+
+	t.Run("different timestamps produce different headers", func(t *testing.T) {
+		mockSigner := &MockSigner{}
+
+		ts1 := time.Unix(1000, 0)
+		ts2 := time.Unix(2000, 0)
+
+		mockSigner.
+			On("Sign", t.Context(), hex.EncodeToString(pubKey), mock.Anything).
+			Return([]byte("signature"), nil)
+
+		headers1, err := beholder.NewAuthHeaderV2(t.Context(), pubKey, mockSigner, ts1)
+		require.NoError(t, err)
+
+		headers2, err := beholder.NewAuthHeaderV2(t.Context(), pubKey, mockSigner, ts2)
+		require.NoError(t, err)
+
+		// Headers should be different due to different timestamps
+		assert.NotEqual(t, headers1["X-Beholder-Node-Auth-Token"], headers2["X-Beholder-Node-Auth-Token"])
+
+		mockSigner.AssertExpectations(t)
+	})
+}
+
 func TestStaticAuthHeaderProvider(t *testing.T) {
 	// Create test headers
 	testHeaders := map[string]string{