gitlab.com/infor-cloud/martian-cloud/tharsis/go-limiter@v0.0.0-20230411193226-3247984d5abc/httplimit/httplimit_test.go (about)

     1  package httplimit_test
     2  
     3  import (
     4  	"crypto/sha512"
     5  	"encoding/base64"
     6  	"fmt"
     7  	"log"
     8  	"net/http"
     9  	"time"
    10  
    11  	"gitlab.com/infor-cloud/martian-cloud/tharsis/go-limiter/httplimit"
    12  	"gitlab.com/infor-cloud/martian-cloud/tharsis/go-limiter/memorystore"
    13  )
    14  
    15  var keyFunc httplimit.KeyFunc
    16  
    17  func ExampleKeyFunc_custom() {
    18  	// This is an example KeyFunc that rate limits using the value from the
    19  	// X-API-Key header. Since this value is likely a secret, it is hashed before
    20  	// passing along to the store.
    21  	keyFunc = httplimit.KeyFunc(func(r *http.Request) (string, error) {
    22  		dig := sha512.Sum512([]byte(r.Header.Get("X-Token")))
    23  		return base64.StdEncoding.EncodeToString(dig[:]), nil
    24  	})
    25  	// middleware, err := httplimit.NewMiddleware(store, keyFunc)
    26  }
    27  
    28  func ExampleIPKeyFunc_headers() {
    29  	keyFunc = httplimit.IPKeyFunc("X-Forwarded-For")
    30  	// middleware, err := httplimit.NewMiddleware(store, keyFunc)
    31  }
    32  
    33  func ExampleNewMiddleware() {
    34  	// Create a store that allows 30 requests per minute.
    35  	store, err := memorystore.New(&memorystore.Config{
    36  		Tokens:   30,
    37  		Interval: time.Minute,
    38  	})
    39  	if err != nil {
    40  		log.Fatal(err)
    41  	}
    42  
    43  	// Create the HTTP middleware from the store, keying by IP address.
    44  	middleware, err := httplimit.NewMiddleware(store, httplimit.IPKeyFunc())
    45  	if err != nil {
    46  		log.Fatal(err)
    47  	}
    48  
    49  	doWork := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    50  		w.WriteHeader(200)
    51  		fmt.Fprintf(w, "hello world")
    52  	})
    53  
    54  	// Wrap an individual handler (only rate limits this endpoint).
    55  	mux1 := http.NewServeMux()
    56  	mux1.Handle("/foo", middleware.Handle(doWork)) // rate limited
    57  	mux1.Handle("/bar", doWork)                    // not rate limited
    58  	_ = mux1
    59  
    60  	// Or wrap the entire mux (rate limits all endpoints).
    61  	mux2 := http.NewServeMux()
    62  	mux2.Handle("/foo", doWork)
    63  	mux2.Handle("/bar", doWork)
    64  	router := middleware.Handle(mux2) // all endpoints are rate limited
    65  	_ = router
    66  }