github.com/letsencrypt/boulder@v0.20251208.0/ratelimits/gcra_test.go (about)

     1  package ratelimits
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/jmhodges/clock"
     8  
     9  	"github.com/letsencrypt/boulder/config"
    10  	"github.com/letsencrypt/boulder/test"
    11  )
    12  
    13  func TestDecide(t *testing.T) {
    14  	clk := clock.NewFake()
    15  	limit := &Limit{Burst: 10, Count: 1, Period: config.Duration{Duration: time.Second}}
    16  	limit.precompute()
    17  
    18  	// Begin by using 1 of our 10 requests.
    19  	d := maybeSpend(clk, Transaction{"test", limit, 1, true, true, false}, clk.Now())
    20  	test.Assert(t, d.allowed, "should be allowed")
    21  	test.AssertEquals(t, d.remaining, int64(9))
    22  	test.AssertEquals(t, d.retryIn, time.Duration(0))
    23  	test.AssertEquals(t, d.resetIn, time.Second)
    24  	// Transaction is set when we're allowed.
    25  	test.AssertEquals(t, d.transaction, Transaction{"test", limit, 1, true, true, false})
    26  
    27  	// Immediately use another 9 of our remaining requests.
    28  	d = maybeSpend(clk, Transaction{"test", limit, 9, true, true, false}, d.newTAT)
    29  	test.Assert(t, d.allowed, "should be allowed")
    30  	test.AssertEquals(t, d.remaining, int64(0))
    31  	// We should have to wait 1 second before we can use another request but we
    32  	// used 9 so we should have to wait 9 seconds to make an identical request.
    33  	test.AssertEquals(t, d.retryIn, time.Second*9)
    34  	test.AssertEquals(t, d.resetIn, time.Second*10)
    35  
    36  	// Our new TAT should be 10 seconds (limit.Burst) in the future.
    37  	test.AssertEquals(t, d.newTAT, clk.Now().Add(time.Second*10))
    38  
    39  	// Let's try using just 1 more request without waiting.
    40  	d = maybeSpend(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
    41  	test.Assert(t, !d.allowed, "should not be allowed")
    42  	test.AssertEquals(t, d.remaining, int64(0))
    43  	test.AssertEquals(t, d.retryIn, time.Second)
    44  	test.AssertEquals(t, d.resetIn, time.Second*10)
    45  	// Transaction is set when we're denied.
    46  	test.AssertEquals(t, d.transaction, Transaction{"test", limit, 1, true, true, false})
    47  
    48  	// Let's try being exactly as patient as we're told to be.
    49  	clk.Add(d.retryIn)
    50  	d = maybeSpend(clk, Transaction{"test", limit, 0, true, true, false}, d.newTAT)
    51  	test.AssertEquals(t, d.remaining, int64(1))
    52  
    53  	// We are 1 second in the future, we should have 1 new request.
    54  	d = maybeSpend(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
    55  	test.Assert(t, d.allowed, "should be allowed")
    56  	test.AssertEquals(t, d.remaining, int64(0))
    57  	test.AssertEquals(t, d.retryIn, time.Second)
    58  	test.AssertEquals(t, d.resetIn, time.Second*10)
    59  
    60  	// Let's try waiting (10 seconds) for our whole bucket to refill.
    61  	clk.Add(d.resetIn)
    62  
    63  	// We should have 10 new requests. If we use 1 we should have 9 remaining.
    64  	d = maybeSpend(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
    65  	test.Assert(t, d.allowed, "should be allowed")
    66  	test.AssertEquals(t, d.remaining, int64(9))
    67  	test.AssertEquals(t, d.retryIn, time.Duration(0))
    68  	test.AssertEquals(t, d.resetIn, time.Second)
    69  
    70  	// Wait just shy of how long we're told to wait for refilling.
    71  	clk.Add(d.resetIn - time.Millisecond)
    72  
    73  	// We should still have 9 remaining because we're still 1ms shy of the
    74  	// refill time.
    75  	d = maybeSpend(clk, Transaction{"test", limit, 0, true, true, false}, d.newTAT)
    76  	test.Assert(t, d.allowed, "should be allowed")
    77  	test.AssertEquals(t, d.remaining, int64(9))
    78  	test.AssertEquals(t, d.retryIn, time.Duration(0))
    79  	test.AssertEquals(t, d.resetIn, time.Millisecond)
    80  
    81  	// Spending 0 simply informed us that we still have 9 remaining, let's see
    82  	// what we have after waiting 20 hours.
    83  	clk.Add(20 * time.Hour)
    84  
    85  	// C'mon, big money, no whammies, no whammies, STOP!
    86  	d = maybeSpend(clk, Transaction{"test", limit, 0, true, true, false}, d.newTAT)
    87  	test.Assert(t, d.allowed, "should be allowed")
    88  	test.AssertEquals(t, d.remaining, int64(10))
    89  	test.AssertEquals(t, d.retryIn, time.Duration(0))
    90  	test.AssertEquals(t, d.resetIn, time.Duration(0))
    91  
    92  	// Turns out that the most we can accrue is 10 (limit.Burst). Let's empty
    93  	// this bucket out so we can try something else.
    94  	d = maybeSpend(clk, Transaction{"test", limit, 10, true, true, false}, d.newTAT)
    95  	test.Assert(t, d.allowed, "should be allowed")
    96  	test.AssertEquals(t, d.remaining, int64(0))
    97  	// We should have to wait 1 second before we can use another request but we
    98  	// used 10 so we should have to wait 10 seconds to make an identical
    99  	// request.
   100  	test.AssertEquals(t, d.retryIn, time.Second*10)
   101  	test.AssertEquals(t, d.resetIn, time.Second*10)
   102  
   103  	// If you spend 0 while you have 0 you should get 0.
   104  	d = maybeSpend(clk, Transaction{"test", limit, 0, true, true, false}, d.newTAT)
   105  	test.Assert(t, d.allowed, "should be allowed")
   106  	test.AssertEquals(t, d.remaining, int64(0))
   107  	test.AssertEquals(t, d.retryIn, time.Duration(0))
   108  	test.AssertEquals(t, d.resetIn, time.Second*10)
   109  
   110  	// We don't play by the rules, we spend 1 when we have 0.
   111  	d = maybeSpend(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
   112  	test.Assert(t, !d.allowed, "should not be allowed")
   113  	test.AssertEquals(t, d.remaining, int64(0))
   114  	test.AssertEquals(t, d.retryIn, time.Second)
   115  	test.AssertEquals(t, d.resetIn, time.Second*10)
   116  
   117  	// Okay, maybe we should play by the rules if we want to get anywhere.
   118  	clk.Add(d.retryIn)
   119  
   120  	// Our patience pays off, we should have 1 new request. Let's use it.
   121  	d = maybeSpend(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
   122  	test.Assert(t, d.allowed, "should be allowed")
   123  	test.AssertEquals(t, d.remaining, int64(0))
   124  	test.AssertEquals(t, d.retryIn, time.Second)
   125  	test.AssertEquals(t, d.resetIn, time.Second*10)
   126  
   127  	// Refill from empty to 5.
   128  	clk.Add(d.resetIn / 2)
   129  
   130  	// Attempt to spend 7 when we only have 5. We should be denied but the
   131  	// decision should reflect a retry of 2 seconds, the time it would take to
   132  	// refill from 5 to 7.
   133  	d = maybeSpend(clk, Transaction{"test", limit, 7, true, true, false}, d.newTAT)
   134  	test.Assert(t, !d.allowed, "should not be allowed")
   135  	test.AssertEquals(t, d.remaining, int64(5))
   136  	test.AssertEquals(t, d.retryIn, time.Second*2)
   137  	test.AssertEquals(t, d.resetIn, time.Second*5)
   138  }
   139  
   140  func TestMaybeRefund(t *testing.T) {
   141  	clk := clock.NewFake()
   142  	limit := &Limit{Burst: 10, Count: 1, Period: config.Duration{Duration: time.Second}}
   143  	limit.precompute()
   144  
   145  	// Begin by using 1 of our 10 requests.
   146  	d := maybeSpend(clk, Transaction{"test", limit, 1, true, true, false}, clk.Now())
   147  	test.Assert(t, d.allowed, "should be allowed")
   148  	test.AssertEquals(t, d.remaining, int64(9))
   149  	test.AssertEquals(t, d.retryIn, time.Duration(0))
   150  	test.AssertEquals(t, d.resetIn, time.Second)
   151  	// Transaction is set when we're refunding.
   152  	test.AssertEquals(t, d.transaction, Transaction{"test", limit, 1, true, true, false})
   153  
   154  	// Refund back to 10.
   155  	d = maybeRefund(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
   156  	test.AssertEquals(t, d.remaining, int64(10))
   157  	test.AssertEquals(t, d.retryIn, time.Duration(0))
   158  	test.AssertEquals(t, d.resetIn, time.Duration(0))
   159  
   160  	// Refund 0, we should still have 10.
   161  	d = maybeRefund(clk, Transaction{"test", limit, 0, true, true, false}, d.newTAT)
   162  	test.AssertEquals(t, d.remaining, int64(10))
   163  	test.AssertEquals(t, d.retryIn, time.Duration(0))
   164  	test.AssertEquals(t, d.resetIn, time.Duration(0))
   165  
   166  	// Spend 1 more of our 10 requests.
   167  	d = maybeSpend(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
   168  	test.Assert(t, d.allowed, "should be allowed")
   169  	test.AssertEquals(t, d.remaining, int64(9))
   170  	test.AssertEquals(t, d.retryIn, time.Duration(0))
   171  	test.AssertEquals(t, d.resetIn, time.Second)
   172  
   173  	// Wait for our bucket to refill.
   174  	clk.Add(d.resetIn)
   175  
   176  	// Attempt to refund from 10 to 11.
   177  	d = maybeRefund(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
   178  	test.Assert(t, !d.allowed, "should not be allowed")
   179  	test.AssertEquals(t, d.remaining, int64(10))
   180  	test.AssertEquals(t, d.retryIn, time.Duration(0))
   181  	test.AssertEquals(t, d.resetIn, time.Duration(0))
   182  	// Transaction is set when our bucket is full.
   183  	test.AssertEquals(t, d.transaction, Transaction{"test", limit, 1, true, true, false})
   184  
   185  	// Spend 10 all 10 of our requests.
   186  	d = maybeSpend(clk, Transaction{"test", limit, 10, true, true, false}, d.newTAT)
   187  	test.Assert(t, d.allowed, "should be allowed")
   188  	test.AssertEquals(t, d.remaining, int64(0))
   189  	// We should have to wait 1 second before we can use another request but we
   190  	// used 10 so we should have to wait 10 seconds to make an identical
   191  	// request.
   192  	test.AssertEquals(t, d.retryIn, time.Second*10)
   193  	test.AssertEquals(t, d.resetIn, time.Second*10)
   194  
   195  	// Attempt a refund of 10.
   196  	d = maybeRefund(clk, Transaction{"test", limit, 10, true, true, false}, d.newTAT)
   197  	test.AssertEquals(t, d.remaining, int64(10))
   198  	test.AssertEquals(t, d.retryIn, time.Duration(0))
   199  	test.AssertEquals(t, d.resetIn, time.Duration(0))
   200  
   201  	// Wait 11 seconds to catching up to TAT.
   202  	clk.Add(11 * time.Second)
   203  
   204  	// Attempt to refund to 11, then ensure it's still 10.
   205  	d = maybeRefund(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
   206  	test.Assert(t, !d.allowed, "should be allowed")
   207  	test.AssertEquals(t, d.remaining, int64(10))
   208  	test.AssertEquals(t, d.retryIn, time.Duration(0))
   209  	test.AssertEquals(t, d.resetIn, time.Duration(0))
   210  	// Transaction is set when our TAT is in the past.
   211  	test.AssertEquals(t, d.transaction, Transaction{"test", limit, 1, true, true, false})
   212  
   213  	// Spend 5 of our 10 requests, then refund 1.
   214  	d = maybeSpend(clk, Transaction{"test", limit, 5, true, true, false}, d.newTAT)
   215  	d = maybeRefund(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
   216  	test.Assert(t, d.allowed, "should be allowed")
   217  	test.AssertEquals(t, d.remaining, int64(6))
   218  	test.AssertEquals(t, d.retryIn, time.Duration(0))
   219  
   220  	// Wait, a 2.5 seconds to refill to 8.5 requests.
   221  	clk.Add(time.Millisecond * 2500)
   222  
   223  	// Ensure we have 8.5 requests.
   224  	d = maybeSpend(clk, Transaction{"test", limit, 0, true, true, false}, d.newTAT)
   225  	test.Assert(t, d.allowed, "should be allowed")
   226  	test.AssertEquals(t, d.remaining, int64(8))
   227  	test.AssertEquals(t, d.retryIn, time.Duration(0))
   228  	// Check that ResetIn represents the fractional earned request.
   229  	test.AssertEquals(t, d.resetIn, time.Millisecond*1500)
   230  
   231  	// Refund 2 requests, we should only have 10, not 10.5.
   232  	d = maybeRefund(clk, Transaction{"test", limit, 2, true, true, false}, d.newTAT)
   233  	test.AssertEquals(t, d.remaining, int64(10))
   234  	test.AssertEquals(t, d.retryIn, time.Duration(0))
   235  	test.AssertEquals(t, d.resetIn, time.Duration(0))
   236  }