package net

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/gorilla/websocket"
)

// The browser transport previously had no tests. CheckOrigin is the security
// boundary for the /ws upgrade; test it directly, including a loopback peer:
// being on the server host must not re-open the cross-site WebSocket hole for
// a hostile browser.
func TestWebSocketCheckOrigin(t *testing.T) {
	cases := []struct {
		name       string
		origin     string
		host       string
		remoteAddr string
		want       bool
	}{
		{"no origin", "", "mud.example:4000", "203.0.113.5:1234", true},
		{"same host", "http://mud.example:4000", "mud.example:4000", "203.0.113.5:1234", true},
		{"same host https", "https://mud.example:4000", "mud.example:4000", "203.0.113.5:1234", true},
		{"same hostname no port", "http://mud.example", "mud.example:4000", "203.0.113.5:1234", true},
		{"hostile origin", "http://evil.example", "mud.example:4000", "203.0.113.5:1234", false},
		{"loopback hostile origin", "http://evil.example", "mud.example:4000", "127.0.0.1:1234", false},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			req := &http.Request{Header: http.Header{}, Host: tc.host, RemoteAddr: tc.remoteAddr}
			if tc.origin != "" {
				req.Header.Set("Origin", tc.origin)
			}
			if got := upgrader.CheckOrigin(req); got != tc.want {
				t.Fatalf("CheckOrigin(%q, host %q, peer %q) = %v, want %v", tc.origin, tc.host, tc.remoteAddr, got, tc.want)
			}
		})
	}
}

// The /ws upgrade wires a live session and writer. Drive it end-to-end over a
// real HTTP test server (the dialer sends no Origin header, which CheckOrigin
// accepts).
func TestWebSocketUpgradeAccepted(t *testing.T) {
	s := &Server{hub: NewHub()}
	done := make(chan struct{})
	s.handler = func(*Session, string) { <-done }
	srv := httptest.NewServer(newHTTPMux(s))
	defer srv.Close()
	defer close(done)

	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws"
	conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
	if err != nil {
		t.Fatalf("upgrade rejected: %v (resp %v)", err, resp)
	}
	conn.Close()
}