From 996610f021ff45fdc98c2ce7884d5fa4e7f9199b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gabriel=20P=C3=A9rez=20S?= <gabriel@guero.net>
Date: Sun, 26 Mar 2017 23:29:12 +0200
Subject: [PATCH] Schema changes detection with HostClient

Related with the issue https://github.com/valyala/fasthttp/issues/244
---
 client.go      |  25 ++++++++
 client_test.go | 152 +++++++++++++++++++++++++++++++++++++++++++++++++
 http.go        |   3 +
 3 files changed, 180 insertions(+)

diff --git a/client.go b/client.go
index 64be717ada..6d7efe63bf 100644
--- a/client.go
+++ b/client.go
@@ -770,9 +770,24 @@ func doRequestFollowRedirects(req *Request, dst []byte, url string, c clientDoer
 	resp.keepBodyBuffer = true
 	oldBody := bodyBuf.B
 	bodyBuf.B = dst
+	scheme := req.uri.Scheme()
+	req.schemaUpdate = false
 
 	redirectsCount := 0
 	for {
+		// In case redirect to different scheme
+		if redirectsCount > 0 && !bytes.Equal(scheme, req.uri.Scheme()) {
+			if strings.HasPrefix(url, string(strHTTPS)) {
+				req.isTLS = true
+				req.uri.SetSchemeBytes(strHTTPS)
+			} else {
+				req.isTLS = false
+				req.uri.SetSchemeBytes(strHTTP)
+			}
+			scheme = req.uri.Scheme()
+			req.schemaUpdate = true
+		}
+
 		req.parsedURI = false
 		req.Header.host = req.Header.host[:0]
 		req.SetRequestURI(url)
@@ -1068,6 +1083,16 @@ func (c *HostClient) doNonNilReqResp(req *Request, resp *Response) (bool, error)
 	// so the GC may reclaim these resources (e.g. response body).
 	resp.Reset()
 
+	// If we detected a redirect to another schema
+	if req.schemaUpdate {
+		c.IsTLS = bytes.Equal(req.URI().Scheme(), strHTTPS)
+		c.Addr = addMissingPort(string(req.Host()), c.IsTLS)
+		c.addrIdx = 0
+		c.addrs = nil
+		req.schemaUpdate = false
+		req.SetConnectionClose()
+	}
+
 	cc, err := c.acquireConn()
 	if err != nil {
 		return false, err
diff --git a/client_test.go b/client_test.go
index 614af92c65..45c54317e2 100644
--- a/client_test.go
+++ b/client_test.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"net/url"
 	"os"
 	"runtime"
 	"strings"
@@ -49,6 +50,157 @@ func TestClientPostArgs(t *testing.T) {
 	}
 }
 
+func TestClientRedirectSameSchema(t *testing.T) {
+
+	listenHTTPS1 := testClientRedirectListener(t, true)
+	defer listenHTTPS1.Close()
+
+	listenHTTPS2 := testClientRedirectListener(t, true)
+	defer listenHTTPS2.Close()
+
+	sHTTPS1 := testClientRedirectChangingSchemaServer(t, listenHTTPS1, listenHTTPS1, true)
+	defer sHTTPS1.Stop()
+
+	sHTTPS2 := testClientRedirectChangingSchemaServer(t, listenHTTPS2, listenHTTPS2, false)
+	defer sHTTPS2.Stop()
+
+	destURL := fmt.Sprintf("https://%s/baz", listenHTTPS1.Addr().String())
+
+	urlParsed, err := url.Parse(destURL)
+	if err != nil {
+		fmt.Println(err)
+		return
+	}
+
+	var reqClient *HostClient
+
+	reqClient = &HostClient{
+		IsTLS: true,
+		Addr:  urlParsed.Host,
+		TLSConfig: &tls.Config{
+			InsecureSkipVerify: true,
+		},
+	}
+
+	statusCode, _, err := reqClient.GetTimeout(nil, destURL, 4000*time.Millisecond)
+	if err != nil {
+		t.Fatalf("HostClient error: %s", err)
+		return
+	}
+
+	if statusCode != 200 {
+		t.Fatalf("HostClient error code response %d", statusCode)
+		return
+	}
+
+}
+
+func TestClientRedirectChangingSchemaHttp2Https(t *testing.T) {
+
+	listenHTTPS := testClientRedirectListener(t, true)
+	defer listenHTTPS.Close()
+
+	listenHTTP := testClientRedirectListener(t, false)
+	defer listenHTTP.Close()
+
+	sHTTPS := testClientRedirectChangingSchemaServer(t, listenHTTPS, listenHTTP, true)
+	defer sHTTPS.Stop()
+
+	sHTTP := testClientRedirectChangingSchemaServer(t, listenHTTPS, listenHTTP, false)
+	defer sHTTP.Stop()
+
+	destURL := fmt.Sprintf("http://%s/baz", listenHTTP.Addr().String())
+
+	urlParsed, err := url.Parse(destURL)
+	if err != nil {
+		fmt.Println(err)
+		return
+	}
+
+	var reqClient *HostClient
+
+	reqClient = &HostClient{
+		Addr: urlParsed.Host,
+		TLSConfig: &tls.Config{
+			InsecureSkipVerify: true,
+		},
+	}
+
+	statusCode, _, err := reqClient.GetTimeout(nil, destURL, 4000*time.Millisecond)
+	if err != nil {
+		t.Fatalf("HostClient error: %s", err)
+		return
+	}
+
+	if statusCode != 200 {
+		t.Fatalf("HostClient error code response %d", statusCode)
+		return
+	}
+
+}
+
+func testClientRedirectListener(t *testing.T, isTLS bool) net.Listener {
+
+	var ln net.Listener
+	var err error
+	var tlsConfig *tls.Config
+
+	if isTLS {
+		certFile := "./ssl-cert-snakeoil.pem"
+		keyFile := "./ssl-cert-snakeoil.key"
+		cert, err1 := tls.LoadX509KeyPair(certFile, keyFile)
+		if err1 != nil {
+			t.Fatalf("Cannot load TLS certificate: %s", err1)
+		}
+		tlsConfig = &tls.Config{
+			Certificates: []tls.Certificate{cert},
+		}
+		ln, err = tls.Listen("tcp", "localhost:0", tlsConfig)
+	} else {
+		ln, err = net.Listen("tcp", "localhost:0")
+	}
+
+	if err != nil {
+		t.Fatalf("cannot listen isTLS %v: %s", isTLS, err)
+	}
+
+	return ln
+}
+
+func testClientRedirectChangingSchemaServer(t *testing.T, https, http net.Listener, isTLS bool) *testEchoServer {
+	s := &Server{
+		Handler: func(ctx *RequestCtx) {
+			if ctx.IsTLS() {
+				ctx.SetStatusCode(200)
+			} else {
+				ctx.Redirect(fmt.Sprintf("https://%s/baz", https.Addr().String()), 301)
+			}
+		},
+	}
+
+	var ln net.Listener
+	if isTLS {
+		ln = https
+	} else {
+		ln = http
+	}
+
+	ch := make(chan struct{})
+	go func() {
+		err := s.Serve(ln)
+		if err != nil {
+			t.Fatalf("unexpected error returned from Serve(): %s", err)
+		}
+		close(ch)
+	}()
+	return &testEchoServer{
+		s:  s,
+		ln: ln,
+		ch: ch,
+		t:  t,
+	}
+}
+
 func TestClientHeaderCase(t *testing.T) {
 	ln := fasthttputil.NewInmemoryListener()
 	defer ln.Close()
diff --git a/http.go b/http.go
index 0ee78a2e7c..6bae63cde2 100644
--- a/http.go
+++ b/http.go
@@ -44,6 +44,9 @@ type Request struct {
 	keepBodyBuffer bool
 
 	isTLS bool
+
+	// To detect scheme changes in redirects
+	schemaUpdate bool
 }
 
 // Response represents HTTP response.