Golang route.SetTable函数代码示例

func watchBackend() {
	var (
		last   string
		svccfg string
		mancfg string
	)

	svc := registry.Default.WatchServices()
	man := registry.Default.WatchManual()

	for {
		select {
		case svccfg = <-svc:
		case mancfg = <-man:
		}

		// manual config overrides service config
		// order matters
		next := svccfg + "\n" + mancfg
		if next == last {
			continue
		}

		t, err := route.ParseString(next)
		if err != nil {
			log.Printf("[WARN] %s", err)
			continue
		}
		route.SetTable(t)

		last = next
	}
}

func TestProxyProducesCorrectXffHeader(t *testing.T) {
	got := "not called"
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got = r.Header.Get("X-Forwarded-For")
	}))
	defer server.Close()

	route.SetTable(mustParseTable("route add mock / " + server.URL))

	tr := &http.Transport{Dial: (&net.Dialer{}).Dial}
	proxy := NewHTTPProxy(tr, config.Proxy{LocalIP: "1.1.1.1", ClientIPHeader: "X-Forwarded-For"})

	req := &http.Request{
		RequestURI: "/",
		Header:     http.Header{"X-Forwarded-For": {"3.3.3.3"}},
		RemoteAddr: "2.2.2.2:666",
		URL:        &url.URL{},
	}

	proxy.ServeHTTP(httptest.NewRecorder(), req)

	if want := "3.3.3.3, 2.2.2.2"; got != want {
		t.Errorf("got %v, but want %v", got, want)
	}
}

func (w *Watcher) Watch() {
	var (
		auto   []string
		manual []string
		t      route.Table
		err    error

		autoConfig   = make(chan []string)
		manualConfig = make(chan []string)
	)

	go watchAutoConfig(w.client, w.tagPrefix, autoConfig)
	go watchManualConfig(w.client, w.configPath, manualConfig)

	for {
		select {
		case auto = <-autoConfig:
		case manual = <-manualConfig:
		}

		if len(auto) == 0 && len(manual) == 0 {
			continue
		}

		input := strings.Join(append(auto, manual...), "\n")
		t, err = route.ParseString(input)
		if err != nil {
			log.Printf("[WARN] %s", err)
			continue
		}
		route.SetTable(t)
	}
}

func initDynamicRoutes() {
	go func() {
		ch := be.Watch()
		for {
			r := <-ch
			t, err := route.ParseString(r)
			if err != nil {
				log.Printf("[WARN] %s", err)
				continue
			}
			route.SetTable(t)
		}
	}()
}

func TestProxyNoRouteStaus(t *testing.T) {
	route.SetTable(make(route.Table))
	tr := &http.Transport{Dial: (&net.Dialer{}).Dial}
	cfg := config.Proxy{NoRouteStatus: 999}
	proxy := NewHTTPProxy(tr, cfg)
	req := &http.Request{
		RequestURI: "/",
		URL:        &url.URL{},
	}

	rec := httptest.NewRecorder()
	proxy.ServeHTTP(rec, req)
	if got, want := rec.Code, cfg.NoRouteStatus; got != want {
		t.Fatalf("got %d want %d", got, want)
	}
}

func initStaticRoutes(routes string) {
	var err error
	var t route.Table

	if strings.HasPrefix(routes, "@") {
		routes = routes[1:]
		log.Print("[INFO] Using static routes from ", routes)
		t, err = route.ParseFile(routes)
	} else {
		log.Print("[INFO] Using static routes from config file")
		t, err = route.ParseString(routes)
	}

	if err != nil {
		log.Fatal("[FATAL] ", err)
	}

	route.SetTable(t)
}

func TestGracefulShutdown(t *testing.T) {
	req := func(url string) int {
		resp, err := http.Get(url)
		if err != nil {
			t.Fatal(err)
		}
		defer resp.Body.Close()
		return resp.StatusCode
	}

	// start a server which responds after the shutdown has been triggered.
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		<-quit // wait for shutdown signal
		return
	}))
	defer srv.Close()

	// load the routing table
	tbl, err := route.ParseString("route add svc / " + srv.URL)
	if err != nil {
		t.Fatal(err)
	}
	route.SetTable(tbl)

	// start proxy with graceful shutdown period long enough
	// to complete one more request.
	var wg sync.WaitGroup
	l := config.Listen{Addr: "127.0.0.1:57777"}
	wg.Add(1)
	go func() {
		defer wg.Done()
		startListeners([]config.Listen{l}, 250*time.Millisecond, route.NewProxy(http.DefaultTransport, config.Proxy{}))
	}()

	// trigger shutdown after some time
	shutdownDelay := 100 * time.Millisecond
	go func() {
		time.Sleep(shutdownDelay)
		close(quit)
	}()

	// give proxy some time to start up
	// needs to be done before shutdown is triggered
	time.Sleep(shutdownDelay / 2)

	// make 200 OK request
	// start before and complete after shutdown was triggered
	if got, want := req("http://"+l.Addr+"/"), 200; got != want {
		t.Fatalf("request 1: got %v want %v", got, want)
	}

	// make 503 request
	// start and complete after shutdown was triggered
	if got, want := req("http://"+l.Addr+"/"), 503; got != want {
		t.Fatalf("got %v want %v", got, want)
	}

	// wait for listen() to return
	// note that the actual listeners have not returned yet
	wg.Wait()
}

func TestProxyGzipHandler(t *testing.T) {
	tests := []struct {
		desc            string
		content         http.HandlerFunc
		acceptEncoding  string
		contentEncoding string
		wantResponse    []byte
	}{
		{
			desc:            "plain body - compressed response",
			content:         plainHandler("text/plain"),
			acceptEncoding:  "gzip",
			contentEncoding: "gzip",
			wantResponse:    gzipContent,
		},
		{
			desc:            "plain body - compressed response (with charset)",
			content:         plainHandler("text/plain; charset=UTF-8"),
			acceptEncoding:  "gzip",
			contentEncoding: "gzip",
			wantResponse:    gzipContent,
		},
		{
			desc:            "compressed body - compressed response",
			content:         gzipHandler("text/plain; charset=UTF-8"),
			acceptEncoding:  "gzip",
			contentEncoding: "gzip",
			wantResponse:    gzipContent,
		},
		{
			desc:            "plain body - plain response",
			content:         plainHandler("text/plain"),
			acceptEncoding:  "",
			contentEncoding: "",
			wantResponse:    plainContent,
		},
		{
			desc:            "compressed body - plain response",
			content:         gzipHandler("text/plain"),
			acceptEncoding:  "",
			contentEncoding: "",
			wantResponse:    plainContent,
		},
		{
			desc:            "plain body - plain response (no match)",
			content:         plainHandler("text/javascript"),
			acceptEncoding:  "gzip",
			contentEncoding: "",
			wantResponse:    plainContent,
		},
	}

	for _, tt := range tests {
		tt := tt // capture loop var
		t.Run(tt.desc, func(t *testing.T) {
			server := httptest.NewServer(tt.content)
			defer server.Close()

			route.SetTable(mustParseTable("route add mock / " + server.URL))

			tr := &http.Transport{Dial: (&net.Dialer{}).Dial}
			proxy := NewHTTPProxy(tr, config.Proxy{GZIPContentTypes: regexp.MustCompile("^text/plain(;.*)?$")})
			req := &http.Request{RequestURI: "/", RemoteAddr: "2.2.2.2:2222", Header: http.Header{"Accept-Encoding": []string{tt.acceptEncoding}}, URL: &url.URL{}}
			rec := httptest.NewRecorder()
			proxy.ServeHTTP(rec, req)

			if got, want := rec.Code, 200; got != want {
				t.Fatalf("got code %d want %d", got, want)
			}
			if got, want := rec.Header().Get("Content-Encoding"), tt.contentEncoding; got != want {
				t.Errorf("got content-encoding %q want %q", got, want)
			}
			if got, want := rec.Body.Bytes(), tt.wantResponse; !bytes.Equal(got, want) {
				t.Errorf("got body %q want %q", got, want)
			}
		})
	}
}