tp: keep polling until the caller's deadline

tp/client.go

@@ -295,7 +295,7 @@ pollLoop:
 			bo = c.nextBO(bo)
 
 			select {
-			case <-time.After(bo):
+			case <-time.After(boBeforeDeadline(ctx, bo)):
 				continue pollLoop
 			case <-ctx.Done():
 				return "", ctx.Err()
@@ -335,10 +335,7 @@ func (c *Client) nextBO(lastBO time.Duration) time.Duration {
 	if c.pollBackoffNext != nil {
 		return c.pollBackoffNext(lastBO)
 	}
-	if lastBO == 0 {
-		return time.Second
-	}
-	return 2 * lastBO
+	return defaultBackoff(lastBO)
 }
 
 func (c *Client) openUserInteractiveURL(ctx context.Context, url string) error {
@@ -392,3 +389,50 @@ func defaultBackoff(lastBO time.Duration) (nextBO time.Duration) {
 	}
 	return 2 * lastBO
 }
+
+const (
+	// pollDeadlineSlack is how much room we try to leave between the last poll
+	// request and the context's deadline, so the request has time to finish.
+	pollDeadlineSlack = time.Second
+
+	// minPollBO is the shortest interval we'll clamp a backoff down to. If
+	// there isn't room for another poll and this much waiting, we let the
+	// context expire rather than spinning on the third party.
+	minPollBO = 100 * time.Millisecond
+)
+
+// boBeforeDeadline returns how long to actually wait before the next poll.
+//
+// Backoffs grow, deadlines don't. A backoff that has grown past what's left of
+// ctx puts the client to sleep for the remainder of its own budget: it stops
+// asking the third party well before it's out of time, and then reports a
+// timeout for a discharge that may have been ready for most of that window.
+// When the next backoff would sleep through the deadline, wait just long
+// enough that one final poll still lands before it.
+//
+// The returned duration is never longer than bo, and the caller keeps
+// backing off from the unclamped value, so this only ever adds a poll.
+func boBeforeDeadline(ctx context.Context, bo time.Duration) time.Duration {
+	deadline, ok := ctx.Deadline()
+	if !ok {
+		return bo
+	}
+
+	remaining := time.Until(deadline)
+
+	slack := remaining / 2
+	if slack > pollDeadlineSlack {
+		slack = pollDeadlineSlack
+	}
+
+	switch latest := remaining - slack; {
+	case latest >= bo:
+		// the backoff already lands before the deadline.
+		return bo
+	case latest >= minPollBO:
+		return latest
+	default:
+		// no room left for another poll.
+		return bo
+	}
+}

tp/client_test.go

@@ -1,7 +1,9 @@
 package tp
 
 import (
+	"context"
 	"testing"
+	"time"
 
 	"github.com/alecthomas/assert/v2"
 	"github.com/hashicorp/go-cleanhttp"
@@ -16,3 +18,43 @@ func TestClient(t *testing.T) {
 	assert.Equal(t, "bar", c1.http.Transport.(*authenticatedHTTP).auth["foo"])
 	assert.Equal(t, "baz", c2.http.Transport.(*authenticatedHTTP).auth["foo"])
 }
+
+func TestBackoffBeforeDeadline(t *testing.T) {
+	t.Run("no deadline", func(t *testing.T) {
+		assert.Equal(t, time.Hour, boBeforeDeadline(context.Background(), time.Hour))
+	})
+
+	t.Run("backoff lands before deadline", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+		defer cancel()
+
+		assert.Equal(t, time.Second, boBeforeDeadline(ctx, time.Second))
+	})
+
+	t.Run("backoff sleeps past deadline", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+
+		// clamped to just under the remaining time, so a final poll lands
+		// before the deadline instead of after it.
+		bo := boBeforeDeadline(ctx, time.Hour)
+		assert.True(t, bo > 8*time.Second, "got %s", bo)
+		assert.True(t, bo < 9*time.Second, "got %s", bo)
+	})
+
+	t.Run("deadline too close for another poll", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
+		defer cancel()
+
+		// nothing useful to clamp to: leave the backoff alone and let the
+		// context expire.
+		assert.Equal(t, time.Hour, boBeforeDeadline(ctx, time.Hour))
+	})
+
+	t.Run("expired deadline", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(context.Background(), 0)
+		defer cancel()
+
+		assert.Equal(t, time.Hour, boBeforeDeadline(ctx, time.Hour))
+	})
+}

tp/tp_test.go

@@ -152,6 +152,56 @@ func TestTP(t *testing.T) {
 		assert.Equal(t, []string{"fp-cav", "dis-cav"}, cavs)
 	})
 
+	t.Run("poll response with backoff past deadline", func(t *testing.T) {
+		pollSecret := ""
+		pollSecretSet := make(chan struct{})
+
+		handleInit = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			_, err := CaveatsFromRequest(r)
+			assert.NoError(t, err)
+
+			pollSecret = tp.RespondPoll(w, r)
+			close(pollSecretSet)
+		})
+
+		hdr := genFP(t, tp, myCaveat("fp-cav"))
+
+		// after the first poll the backoff is far longer than the context
+		// allows for. The client has to notice that and squeeze a final poll in
+		// before the deadline, or it spends the rest of its budget asleep and
+		// never sees a discharge that showed up in the meantime.
+		c := NewClient(firstPartyLocation,
+			WithPollingBackoff(func(last time.Duration) time.Duration {
+				if last == 0 {
+					return 10 * time.Millisecond
+				}
+				return 10 * time.Minute
+			}),
+		)
+
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+
+		go func() {
+			select {
+			case <-pollSecretSet:
+				select {
+				case <-time.After(200 * time.Millisecond):
+					assert.NoError(t, tp.DischargePoll(context.Background(), pollSecret, myCaveat("dis-cav")))
+				case <-ctx.Done():
+					panic("oh no")
+				}
+			case <-ctx.Done():
+				panic("oh no")
+			}
+		}()
+
+		hdr, err = c.FetchDischargeTokens(ctx, hdr)
+		assert.NoError(t, err)
+		cavs := checkFP(t, hdr)
+		assert.Equal(t, []string{"fp-cav", "dis-cav"}, cavs)
+	})
+
 	t.Run("user interactive response", func(t *testing.T) {
 		userSecret := ""