Add unit test for proxy and shutdown channel

* Proxy has initial unit test and more can be added
* Shutdown channel and cancellation added for proper shutdown of
the proxy

Signed-off-by: Alex Ellis (OpenFaaS Ltd) <alexellis2@gmail.com>

cmd/up.go

@@ -77,6 +77,7 @@ func runUp(_ *cobra.Command, _ []string) error {
 
 	shutdownTimeout := time.Second * 1
 	timeout := time.Second * 60
+	proxyDoneCh := make(chan bool)
 
 	wg := sync.WaitGroup{}
 	wg.Add(1)
@@ -92,14 +93,18 @@ func runUp(_ *cobra.Command, _ []string) error {
 		if err != nil {
 			fmt.Println(err)
 		}
+
+		// Close proxy
+		proxyDoneCh <- true
 		time.AfterFunc(shutdownTimeout, func() {
 			wg.Done()
 		})
 	}()
 
 	gatewayURLChan := make(chan string, 1)
-	proxy := pkg.NewProxy(timeout)
+	proxyPort := 8080
-	go proxy.Start(gatewayURLChan)
+	proxy := pkg.NewProxy(proxyPort, timeout)
+	go proxy.Start(gatewayURLChan, proxyDoneCh)
 
 	go func() {
 		wd, _ := os.Getwd()
@@ -119,7 +124,7 @@ func runUp(_ *cobra.Command, _ []string) error {
 			}
 		}
 		log.Printf("[up] Sending %s to proxy\n", host)
-		gatewayURLChan <- host
+		gatewayURLChan <- host + ":8080"
 		close(gatewayURLChan)
 	}()
 

pkg/proxy.go

@@ -1,6 +1,7 @@
 package pkg
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -10,47 +11,82 @@ import (
 	"time"
 )
 
-func NewProxy(timeout time.Duration) *Proxy {
+func NewProxy(port int, timeout time.Duration) *Proxy {
 
 	return &Proxy{
+		Port:    port,
 		Timeout: timeout,
 	}
 }
 
 type Proxy struct {
 	Timeout time.Duration
+	Port    int
 }
 
-func (p *Proxy) Start(gatewayChan chan string) error {
+func (p *Proxy) Start(gatewayChan chan string, done chan bool) error {
-	tcp := 8080
+	tcp := p.Port
 
 	http.DefaultClient.CheckRedirect = func(req *http.Request, via []*http.Request) error {
 		return http.ErrUseLastResponse
 	}
-
+	ps := proxyState{
-	data := struct{ host string }{
+		Host: "",
-		host: "",
 	}
 
-	data.host = <-gatewayChan
+	ps.Host = <-gatewayChan
 
 	log.Printf("Starting faasd proxy on %d\n", tcp)
 
-	fmt.Printf("Gateway: %s\n", data.host)
+	fmt.Printf("Gateway: %s\n", ps.Host)
 
 	s := &http.Server{
 		Addr:           fmt.Sprintf(":%d", tcp),
 		ReadTimeout:    p.Timeout,
 		WriteTimeout:   p.Timeout,
 		MaxHeaderBytes: 1 << 20, // Max header of 1MB
-		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		Handler:        http.HandlerFunc(makeProxy(&ps)),
+	}
+
+	go func() {
+		log.Printf("[proxy] Begin listen on %d\n", p.Port)
+		if err := s.ListenAndServe(); err != http.ErrServerClosed {
+			log.Printf("Error ListenAndServe: %v", err)
+		}
+	}()
+
+	log.Println("[proxy] Wait for done")
+	<-done
+	log.Println("[proxy] Done received")
+	if err := s.Shutdown(context.Background()); err != nil {
+		log.Printf("[proxy] Error in Shutdown: %v", err)
+	}
+
+	return nil
+}
+
+// copyHeaders clones the header values from the source into the destination.
+func copyHeaders(destination http.Header, source *http.Header) {
+	for k, v := range *source {
+		vClone := make([]string, len(v))
+		copy(vClone, v)
+		destination[k] = vClone
+	}
+}
+
+type proxyState struct {
+	Host string
+}
+
+func makeProxy(ps *proxyState) func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
 
 		query := ""
 		if len(r.URL.RawQuery) > 0 {
 			query = "?" + r.URL.RawQuery
 		}
 
-			upstream := fmt.Sprintf("http://%s:8080%s%s", data.host, r.URL.Path, query)
+		upstream := fmt.Sprintf("http://%s%s%s", ps.Host, r.URL.Path, query)
 		fmt.Printf("[faasd] proxy: %s\n", upstream)
 
 		if r.Body != nil {
@@ -82,18 +118,5 @@ func (p *Proxy) Start(gatewayChan chan string) error {
 
 		w.WriteHeader(upRes.StatusCode)
 		io.Copy(w, upRes.Body)
-
-		}),
-	}
-
-	return s.ListenAndServe()
-}
-
-// copyHeaders clones the header values from the source into the destination.
-func copyHeaders(destination http.Header, source *http.Header) {
-	for k, v := range *source {
-		vClone := make([]string, len(v))
-		copy(vClone, v)
-		destination[k] = vClone
 	}
 }

pkg/proxy_test.go

@@ -0,0 +1,73 @@
+package pkg
+
+import (
+	"fmt"
+	"io/ioutil"
+	"log"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"sync"
+	"testing"
+	"time"
+)
+
+func Test_Proxy_ToPrivateServer(t *testing.T) {
+
+	wantBodyText := "OK"
+	wantBody := []byte(wantBodyText)
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+
+		if r.Body != nil {
+			defer r.Body.Close()
+		}
+
+		w.WriteHeader(http.StatusOK)
+		w.Write(wantBody)
+
+	}))
+
+	defer upstream.Close()
+	port := 8080
+	proxy := NewProxy(port, time.Second*1)
+
+	gwChan := make(chan string, 1)
+	doneCh := make(chan bool)
+
+	go proxy.Start(gwChan, doneCh)
+
+	u, _ := url.Parse(upstream.URL)
+	log.Println("Host", u.Host)
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		gwChan <- u.Host
+		wg.Done()
+	}()
+	wg.Wait()
+
+	req, err := http.NewRequest(http.MethodGet, fmt.Sprintf("http://127.0.0.1:%d", port), nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	for i := 1; i < 11; i++ {
+		res, err := http.DefaultClient.Do(req)
+		if err != nil {
+			t.Logf("Try %d, gave error: %s", i, err)
+
+			time.Sleep(time.Millisecond * 100)
+		} else {
+
+			resBody, _ := ioutil.ReadAll(res.Body)
+			if string(resBody) != string(wantBody) {
+				t.Errorf("want %s, but got %s in body", string(wantBody), string(resBody))
+			}
+			break
+		}
+	}
+
+	go func() {
+		doneCh <- true
+	}()
+}