package main

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"net/http/httputil"
	"net/url"
	"strings"

	"golang.org/x/net/http/httpguts"
)

func upgradeType(h http.Header) string {
	if !httpguts.HeaderValuesContainsToken(h["Connection"], "Upgrade") {
		return ""
	}
	return strings.ToLower(h.Get("Upgrade"))
}

func main() {
	reqDone := make(chan struct{})
	backendServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer close(reqDone)
		if upgradeType(r.Header) != "websocket" {
			http.Error(w, "unexpected request", 400)
			return
		}
		c, _, err := w.(http.Hijacker).Hijack()
		if err != nil {
			http.Error(w, fmt.Sprintf("%v", err), 500)
			return
		}

		io.WriteString(c, "HTTP/1.1 101 Switching Protocols\r\nConnection: upgrade\r\nUpgrade: WebSocket\r\n\r\n")

		if tcpc, ok := c.(interface {
			CloseWrite() error
		}); ok {
			tcpc.CloseWrite()
		} else if closer, ok := c.(io.Closer); ok {
			closer.Close()
			return
		}

		bs := bufio.NewScanner(c)
		if !bs.Scan() {
			fmt.Printf("backend failed to read line from client: %v", bs.Err())
			return
		}

		got := bs.Text()
		want := "Hello"
		if got != want {
			panic(fmt.Sprintf("got %#q, want %#q", got, want))
		}
		fmt.Println("backend got", got)
	}))
	defer backendServer.Close()

	backURL, _ := url.Parse(backendServer.URL)
	rproxy := httputil.NewSingleHostReverseProxy(backURL)
	rproxy.ErrorLog = log.New(ioutil.Discard, "", 0) // quiet for tests

	frontendProxy := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		rproxy.ServeHTTP(rw, req)
	}))

	defer frontendProxy.Close()

	req, _ := http.NewRequest("GET", frontendProxy.URL, nil)
	req.Header.Set("Connection", "Upgrade")
	req.Header.Set("Upgrade", "websocket")

	c := frontendProxy.Client()
	res, err := c.Do(req)
	if err != nil {
		panic(err)
	}
	if res.StatusCode != 101 {
		panic(fmt.Sprintf("status = %v; want 101", res.Status))
	}

	if upgradeType(res.Header) != "websocket" {
		panic(fmt.Sprintf("not websocket upgrade; got %#v", res.Header))
	}
	rwc, ok := res.Body.(io.ReadWriteCloser)
	if !ok {
		panic(fmt.Sprintf("response body is of type %T; does not implement ReadWriteCloser", res.Body))
	}
	defer rwc.Close()

	ioutil.ReadAll(rwc)
	io.WriteString(rwc, "Hello\n")
	<-reqDone
}