feat(auth): Add JWT authentication for heartbeat endpoints (#228)

* Add JWT-protected heartbeat endpoint with organization validation

* Fix JWT audience to match token-exchange service

* Fix JWT claims extraction for ESP v1 (GAE Flex)

* Use portable JWT parsing with go-jose library

* Update heartbeat client to use JWT authentication

* feat(auth): add JWT token refresh for heartbeat connections

Implements automatic JWT token renewal before expiry to handle GAE's
1-hour connection limits. Includes comprehensive test coverage.

* update comment

* feat(auth): update handler to read ESPv1's custom header

* add debug logging

* fix(auth): parse claims as a JSON string

* use base64.StdEncoding

* restore expiry parsing from JWT

* Make tests table-driven

Author: robertodauria

cmd/heartbeat/main.go

@@ -1,7 +1,9 @@
 package main
 
 import (
+	"bytes"
 	"context"
+	"encoding/json"
 	"flag"
 	"fmt"
 	"log"
@@ -29,7 +31,7 @@ import (
 
 var (
 	heartbeatURL        string
-	hostname			flagx.StringFile
+	hostname            flagx.StringFile
 	experiment          string
 	pod                 string
 	node                string
@@ -41,13 +43,60 @@ var (
 	heartbeatPeriod     = static.HeartbeatPeriod
 	mainCtx, mainCancel = context.WithCancel(context.Background())
 	lbPath              = "/metadata/loadbalanced"
+
+	// JWT authentication parameters
+	apiKey           string
+	tokenExchangeURL string
 )
 
 // Checker generates a health score for the heartbeat instance (0, 1).
 type Checker interface {
 	GetHealth(ctx context.Context) float64 // Health score.
 }
 
+// TokenResponse represents the response from the token exchange service
+type TokenResponse struct {
+	Token string `json:"token"`
+}
+
+// getJWTTokenFunc is a variable that holds the JWT token function, allowing for test overrides
+var getJWTTokenFunc = getJWTToken
+
+// getJWTToken exchanges an API key for a JWT token
+func getJWTToken(apiKey, tokenExchangeURL string) (string, error) {
+	// Prepare the request payload
+	payload := map[string]string{
+		"api_key": apiKey,
+	}
+	payloadBytes, err := json.Marshal(payload)
+	if err != nil {
+		return "", fmt.Errorf("failed to marshal token request: %w", err)
+	}
+
+	// Make the HTTP request to the token exchange service
+	resp, err := http.Post(tokenExchangeURL, "application/json", bytes.NewBuffer(payloadBytes))
+	if err != nil {
+		return "", fmt.Errorf("failed to request JWT token: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return "", fmt.Errorf("token exchange failed with status %d", resp.StatusCode)
+	}
+
+	// Parse the response
+	var tokenResp TokenResponse
+	if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
+		return "", fmt.Errorf("failed to parse token response: %w", err)
+	}
+
+	if tokenResp.Token == "" {
+		return "", fmt.Errorf("empty token received from exchange service")
+	}
+
+	return tokenResp.Token, nil
+}
+
 func init() {
 	flag.StringVar(&heartbeatURL, "heartbeat-url", "ws://localhost:8080/v2/platform/heartbeat",
 		"URL for locate service")
@@ -59,12 +108,23 @@ func init() {
 	flag.Var(&kubernetesURL, "kubernetes-url", "URL for Kubernetes API")
 	flag.Var(&registrationURL, "registration-url", "URL for site registration")
 	flag.Var(&services, "services", "Maps experiment target names to their set of services")
+	flag.StringVar(&apiKey, "api-key", "", "API key for JWT token exchange (required)")
+	flag.StringVar(&tokenExchangeURL, "token-exchange-url", "https://auth.mlab-sandbox.measurementlab.net/v0/token/autojoin",
+		"URL for token exchange service")
 }
 
 func main() {
 	flag.Parse()
 	rtx.Must(flagx.ArgsFromEnvWithLog(flag.CommandLine, false), "failed to read args from env")
 
+	// Validate JWT authentication parameters
+	if apiKey == "" {
+		log.Fatal("API key is required for JWT authentication (-api-key flag)")
+	}
+	if tokenExchangeURL == "" {
+		log.Fatal("Token exchange URL is required (-token-exchange-url flag)")
+	}
+
 	// Start metrics server.
 	prom := prometheusx.MustServeMetrics()
 	defer prom.Close()
@@ -82,9 +142,26 @@ func main() {
 	rtx.Must(err, "could not load registration data")
 	hbm := v2.HeartbeatMessage{Registration: r}
 
-	// Establish a connection.
+	// Get JWT token for authentication
+	log.Printf("Exchanging API key for JWT token...")
+	jwtToken, err := getJWTTokenFunc(apiKey, tokenExchangeURL)
+	rtx.Must(err, "failed to get JWT token")
+	log.Printf("Successfully obtained JWT token")
+
+	// Prepare headers with JWT authentication
+	headers := http.Header{}
+	headers.Set("Authorization", "Bearer "+jwtToken)
+
+	// Establish a connection with JWT authentication.
 	conn := connection.NewConn()
-	err = conn.Dial(heartbeatURL, http.Header{}, hbm)
+
+	// Set up JWT token refresh for automatic token renewal
+	conn.SetTokenRefresher(func() (string, error) {
+		log.Printf("Refreshing JWT token...")
+		return getJWTTokenFunc(apiKey, tokenExchangeURL)
+	})
+
+	err = conn.Dial(heartbeatURL, headers, hbm)
 	rtx.Must(err, "failed to establish a websocket connection with %s", heartbeatURL)
 
 	probe := health.NewPortProbe(svcs)

cmd/heartbeat/main_test.go

@@ -17,6 +17,14 @@ import (
 )
 
 func Test_main(t *testing.T) {
+	// Override the JWT token function for testing
+	getJWTTokenFunc = func(apiKey, tokenExchangeURL string) (string, error) {
+		return "fake-jwt-token", nil
+	}
+	defer func() {
+		getJWTTokenFunc = getJWTToken // restore original function
+	}()
+
 	mainCtx, mainCancel = context.WithCancel(context.Background())
 	fh := testdata.FakeHandler{}
 	s := testdata.FakeServer(fh.Upgrade)
@@ -36,6 +44,8 @@ func Test_main(t *testing.T) {
 	flag.Set("namespace", "default")
 	flag.Set("registration-url", "file:./registration/testdata/registration.json")
 	flag.Set("services", "ndt/ndt7=ws://:"+u.Port()+"/ndt/v7/download")
+	flag.Set("api-key", "test-api-key")
+	flag.Set("token-exchange-url", "http://fake-token-exchange.example.com/token")
 
 	heartbeatPeriod = 2 * time.Second
 	timer := time.NewTimer(2 * heartbeatPeriod)

connection/connection.go

@@ -7,13 +7,14 @@ import (
 	"log"
 	"net/http"
 	"net/url"
-	"sync"
+	"strings"
 	"time"
 
 	"github.com/cenkalti/backoff/v4"
 	"github.com/gorilla/websocket"
 	"github.com/m-lab/locate/metrics"
 	"github.com/m-lab/locate/static"
+	"gopkg.in/square/go-jose.v2/jwt"
 )
 
 var (
@@ -23,8 +24,14 @@ var (
 	// retryErrors contains the list of errors that may become successful
 	// if the request is retried.
 	retryErrors = map[int]bool{408: true, 425: true, 500: true, 502: true, 503: true, 504: true}
+
+	// JWT refresh buffer: refresh token if it expires within this duration
+	jwtRefreshBuffer = 5 * time.Minute
 )
 
+// TokenRefresher is a function type that can refresh JWT tokens.
+type TokenRefresher func() (string, error)
+
 // Conn contains the state needed to connect, reconnect, and send
 // messages.
 // Default values must be updated before calling `Dial`.
@@ -52,9 +59,11 @@ type Conn struct {
 	url         url.URL
 	header      http.Header
 	ticker      time.Ticker
-	mu          sync.Mutex
 	isDialed    bool
 	isConnected bool
+
+	// JWT token refresh functionality
+	tokenRefresher TokenRefresher
 }
 
 // NewConn creates a new Conn with default values.
@@ -69,6 +78,11 @@ func NewConn() *Conn {
 	return c
 }
 
+// SetTokenRefresher sets the function used to refresh JWT tokens when they expire.
+func (c *Conn) SetTokenRefresher(refresher TokenRefresher) {
+	c.tokenRefresher = refresher
+}
+
 // Dial creates a new persistent client connection and sets
 // the necessary state for future reconnections. It also
 // starts a goroutine to reset the number of reconnections.
@@ -167,6 +181,12 @@ func (c *Conn) close() error {
 // In case of failure, it uses an exponential backoff to
 // increase the duration of retry attempts.
 func (c *Conn) connect() error {
+	// Check if JWT token needs refreshing before attempting connection
+	if err := c.refreshJWTIfNeeded(); err != nil {
+		log.Printf("failed to refresh JWT token: %v", err)
+		// Continue with existing token - it might still work
+	}
+
 	b := c.getBackoff()
 	ticker := backoff.NewTicker(b)
 
@@ -234,3 +254,67 @@ func (c *Conn) getBackoff() *backoff.ExponentialBackOff {
 	b.MaxElapsedTime = c.MaxElapsedTime
 	return b
 }
+
+// parseJWTExpiry extracts the expiry time from a JWT token without verification
+func (c *Conn) parseJWTExpiry(token string) time.Time {
+	parsed, err := jwt.ParseSigned(token)
+	if err != nil {
+		log.Printf("failed to parse JWT for expiry check: %v", err)
+		return time.Time{} // Return zero time on error
+	}
+
+	var claims jwt.Claims
+	if err := parsed.UnsafeClaimsWithoutVerification(&claims); err != nil {
+		log.Printf("failed to extract JWT claims for expiry check: %v", err)
+		return time.Time{} // Return zero time on error
+	}
+
+	if claims.Expiry == nil {
+		return time.Time{} // Return zero time if no expiry
+	}
+
+	return claims.Expiry.Time()
+}
+
+// SetToken sets a JWT token
+func (c *Conn) SetToken(token string) {
+	c.header.Set("Authorization", "Bearer "+token)
+}
+
+func (c *Conn) refreshJWTIfNeeded() error {
+	// Only refresh if we have a token refresher
+	if c.tokenRefresher == nil {
+		return nil
+	}
+
+	// Extract current token from headers
+	authHeader := c.header.Get("Authorization")
+	if !strings.HasPrefix(authHeader, "Bearer ") {
+		return nil // No token present
+	}
+
+	currentToken := strings.TrimPrefix(authHeader, "Bearer ")
+
+	// Parse expiry on-demand
+	expiry := c.parseJWTExpiry(currentToken)
+	if expiry.IsZero() {
+		return nil // Token has no expiry or couldn't parse
+	}
+
+	// Check if refresh needed
+	if time.Until(expiry) > jwtRefreshBuffer {
+		return nil // Still fresh
+	}
+
+	// Token is expired or close to expiring, refresh it
+	log.Printf("JWT token expires soon, refreshing...")
+	newToken, err := c.tokenRefresher()
+	if err != nil {
+		return err
+	}
+
+	c.header.Set("Authorization", "Bearer "+newToken)
+	log.Printf("JWT token refreshed successfully")
+
+	return nil
+}

connection/connection_test.go

@@ -200,6 +200,79 @@ func TestWriteMessage_ClosedServer(t *testing.T) {
 	}
 }
 
+func TestConn_RefreshJWTIfNeeded(t *testing.T) {
+	tests := []struct {
+		name             string
+		token            string
+		hasRefresher     bool
+		refreshedToken   string
+		wantRefreshCount int
+		wantHeaderToken  string
+		wantErr          bool
+	}{
+		{
+			name:             "expired_token_with_refresher",
+			token:            "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjE2MDA5MzUzMDB9.fake-signature", // Expired in 2020
+			hasRefresher:     true,
+			refreshedToken:   "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjk5OTk5OTk5OTl9.fake-signature",
+			wantRefreshCount: 1,
+			wantHeaderToken:  "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjk5OTk5OTk5OTl9.fake-signature",
+		},
+		{
+			name:             "expired_token_no_refresher",
+			token:            "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjE2MDA5MzUzMDB9.fake-signature", // Expired in 2020
+			hasRefresher:     false,
+			wantRefreshCount: 0,
+			wantHeaderToken:  "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjE2MDA5MzUzMDB9.fake-signature",
+		},
+		{
+			name:             "valid_token_with_refresher",
+			token:            "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjk5OTk5OTk5OTl9.fake-signature", // Expires far in future
+			hasRefresher:     true,
+			refreshedToken:   "new-token",
+			wantRefreshCount: 0,
+			wantHeaderToken:  "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjk5OTk5OTk5OTl9.fake-signature",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c := NewConn()
+
+			refreshCount := 0
+			if tt.hasRefresher {
+				c.SetTokenRefresher(func() (string, error) {
+					refreshCount++
+					return tt.refreshedToken, nil
+				})
+			}
+
+			headers := http.Header{}
+			headers.Set("Authorization", "Bearer "+tt.token)
+			c.header = headers
+
+			fh := testdata.FakeHandler{}
+			s := testdata.FakeServer(fh.Upgrade)
+			defer close(c, s)
+
+			err := c.refreshJWTIfNeeded()
+			if (err != nil) != tt.wantErr {
+				t.Errorf("refreshJWTIfNeeded() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+
+			if refreshCount != tt.wantRefreshCount {
+				t.Errorf("Expected %d token refresh calls, got %d", tt.wantRefreshCount, refreshCount)
+			}
+
+			newAuth := c.header.Get("Authorization")
+			if newAuth != tt.wantHeaderToken {
+				t.Errorf("Authorization header = %s, want %s", newAuth, tt.wantHeaderToken)
+			}
+		})
+	}
+}
+
 func close(c *Conn, s *httptest.Server) {
 	c.Close()
 	s.Close()