github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/tools/auth-injector/main.go (about)

     1  package main
     2  
     3  import (
     4  	"flag"
     5  	"log"
     6  	"net"
     7  	"net/http"
     8  	"net/http/httputil"
     9  	"net/url"
    10  	"time"
    11  )
    12  
    13  type Config struct {
    14  	LocalAddress string
    15  	RemoteURL    string
    16  	TenantID     string
    17  }
    18  
    19  func main() {
    20  	cfg := Config{}
    21  	flag.StringVar(&cfg.LocalAddress, "local-address", ":8080", "Local address to listen on (host:port or :port).")
    22  	flag.StringVar(&cfg.RemoteURL, "remote-address", "", "URL of target to forward requests to to (eg. http://domain.com:80).")
    23  	flag.StringVar(&cfg.TenantID, "tenant-id", "", "Tenant ID to inject to proxied requests.")
    24  	flag.Parse()
    25  
    26  	// Parse remote URL.
    27  	if cfg.RemoteURL == "" {
    28  		log.Fatalln("No -remote-address specified.")
    29  	}
    30  
    31  	remoteURL, err := url.Parse(cfg.RemoteURL)
    32  	if err != nil {
    33  		log.Fatalf("Unable to parse remote address. Error: %s.", err.Error())
    34  	}
    35  	log.Println("Forwarding to", remoteURL)
    36  
    37  	ln, err := net.Listen("tcp", cfg.LocalAddress)
    38  	if err != nil {
    39  		log.Fatal(err)
    40  	}
    41  
    42  	s := &http.Server{
    43  		Addr:           cfg.LocalAddress,
    44  		Handler:        injectAuthHeader(cfg.TenantID, httputil.NewSingleHostReverseProxy(remoteURL)),
    45  		ReadTimeout:    10 * time.Second,
    46  		WriteTimeout:   10 * time.Second,
    47  		MaxHeaderBytes: 1 << 20,
    48  	}
    49  
    50  	log.Println("Listening on", ln.Addr())
    51  	log.Fatal(s.Serve(ln))
    52  }
    53  
    54  func injectAuthHeader(tenantID string, h http.Handler) http.Handler {
    55  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    56  		r.Header.Set("X-Scope-OrgID", tenantID)
    57  		h.ServeHTTP(w, r)
    58  	})
    59  }