From 9c1b9a5a74bba8d967bee69b6e3f833d26705688 Mon Sep 17 00:00:00 2001
From: stl0755 <stl0755@gmail.com>
Date: Wed, 14 Aug 2019 15:02:53 +0800
Subject: [PATCH] test

---
 conf/frps_full.ini |   3 ++
 server/service.go  |   9 ++++
 utils/net/wss.go   | 105 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 117 insertions(+)
 create mode 100644 utils/net/wss.go

diff --git a/conf/frps_full.ini b/conf/frps_full.ini
index ed507cef..a780cb2e 100644
--- a/conf/frps_full.ini
+++ b/conf/frps_full.ini
@@ -71,3 +71,6 @@ tcp_mux = true
 
 # custom 404 page for HTTP requests
 # custom_404_page = /path/to/404.html
+
+wss_crt_path = ./wss.crt
+wss_key_path = ./wss.key
diff --git a/server/service.go b/server/service.go
index a4d2e9df..543d0f39 100644
--- a/server/service.go
+++ b/server/service.go
@@ -67,6 +67,9 @@ type Service struct {
 	// Accept connections using websocket
 	websocketListener net.Listener
 
+	// Accept connections using wss
+	wssListener frpNet.Listener
+
 	// Accept frp tls connections
 	tlsListener net.Listener
 
@@ -157,6 +160,12 @@ func NewService(cfg config.ServerCommonConf) (svr *Service, err error) {
 	})
 	svr.websocketListener = frpNet.NewWebsocketListener(websocketLn)
 
+	// frp wss listener
+	wssListener := svr.muxer.Listen(1, 1, func(data []byte) bool {
+		return int(data[0]) == 0x16
+	})
+	svr.wssListener = frpNet.NewWssListener(wssListener)
+
 	// Create http vhost muxer.
 	if cfg.VhostHttpPort > 0 {
 		rp := vhost.NewHttpReverseProxy(vhost.HttpReverseProxyOptions{
diff --git a/utils/net/wss.go b/utils/net/wss.go
new file mode 100644
index 00000000..cf469269
--- /dev/null
+++ b/utils/net/wss.go
@@ -0,0 +1,105 @@
+package net
+
+import (
+	"errors"
+	"fmt"
+	"net"
+	"net/http"
+	"net/url"
+	"time"
+
+	"github.com/fatedier/frp/utils/log"
+
+	"golang.org/x/net/websocket"
+)
+
+var (
+	ErrWssListenerClosed = errors.New("wss listener closed")
+)
+
+type WssListener struct {
+	net.Addr
+	ln     net.Listener
+	accept chan Conn
+	log.Logger
+
+	server    *http.Server
+	httpMutex *http.ServeMux
+}
+
+// NewWssListener to handle wss connections
+// ln: tcp listener for wss connections
+func NewWssListener(ln net.Listener) (wl *WssListener) {
+	wl = &WssListener{
+		Addr:   ln.Addr(),
+		accept: make(chan Conn),
+		Logger: log.NewPrefixLogger(""),
+	}
+
+	muxer := http.NewServeMux()
+	muxer.Handle(FrpWebsocketPath, websocket.Handler(func(c *websocket.Conn) {
+		notifyCh := make(chan struct{})
+		conn := WrapCloseNotifyConn(c, func() {
+			close(notifyCh)
+		})
+		wl.accept <- conn
+		<-notifyCh
+	}))
+
+	wl.server = &http.Server{
+		Addr:    ln.Addr().String(),
+		Handler: muxer,
+	}
+
+	certFile := "a.cert"
+	keyFile := "a.key"
+
+	go wl.server.ServeTLS(ln, certFile, keyFile)
+	return
+}
+
+func ListenWss(bindAddr string, bindPort int) (*WssListener, error) {
+	tcpLn, err := net.Listen("tcp", fmt.Sprintf("%s:%d", bindAddr, bindPort))
+	if err != nil {
+		return nil, err
+	}
+	l := NewWssListener(tcpLn)
+	return l, nil
+}
+
+func (p *WssListener) Accept() (Conn, error) {
+	c, ok := <-p.accept
+	if !ok {
+		return nil, ErrWssListenerClosed
+	}
+	return c, nil
+}
+
+func (p *WssListener) Close() error {
+	return p.server.Close()
+}
+
+// addr: domain:port
+func ConnectWssServer(addr string, httpProtocol string, wsProtocol string) (Conn, error) {
+	addr = wsProtocol + "://" + addr + FrpWebsocketPath
+	uri, err := url.Parse(addr)
+	if err != nil {
+		return nil, err
+	}
+
+	origin := httpProtocol + "://" + uri.Host
+	cfg, err := websocket.NewConfig(addr, origin)
+	if err != nil {
+		return nil, err
+	}
+	cfg.Dialer = &net.Dialer{
+		Timeout: 10 * time.Second,
+	}
+
+	conn, err := websocket.DialConfig(cfg)
+	if err != nil {
+		return nil, err
+	}
+	c := WrapConn(conn)
+	return c, nil
+}