rename models to pkg (#2005)

pkg/util/net/conn.go

@@ -0,0 +1,243 @@
+// Copyright 2016 fatedier, fatedier@gmail.com
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package net
+
+import (
+	"context"
+	"crypto/tls"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"sync/atomic"
+	"time"
+
+	"github.com/fatedier/frp/pkg/util/xlog"
+
+	gnet "github.com/fatedier/golib/net"
+	kcp "github.com/fatedier/kcp-go"
+)
+
+type ContextGetter interface {
+	Context() context.Context
+}
+
+type ContextSetter interface {
+	WithContext(ctx context.Context)
+}
+
+func NewLogFromConn(conn net.Conn) *xlog.Logger {
+	if c, ok := conn.(ContextGetter); ok {
+		return xlog.FromContextSafe(c.Context())
+	}
+	return xlog.New()
+}
+
+func NewContextFromConn(conn net.Conn) context.Context {
+	if c, ok := conn.(ContextGetter); ok {
+		return c.Context()
+	}
+	return context.Background()
+}
+
+// ContextConn is the connection with context
+type ContextConn struct {
+	net.Conn
+
+	ctx context.Context
+}
+
+func NewContextConn(ctx context.Context, c net.Conn) *ContextConn {
+	return &ContextConn{
+		Conn: c,
+		ctx:  ctx,
+	}
+}
+
+func (c *ContextConn) WithContext(ctx context.Context) {
+	c.ctx = ctx
+}
+
+func (c *ContextConn) Context() context.Context {
+	return c.ctx
+}
+
+type WrapReadWriteCloserConn struct {
+	io.ReadWriteCloser
+
+	underConn net.Conn
+}
+
+func WrapReadWriteCloserToConn(rwc io.ReadWriteCloser, underConn net.Conn) net.Conn {
+	return &WrapReadWriteCloserConn{
+		ReadWriteCloser: rwc,
+		underConn:       underConn,
+	}
+}
+
+func (conn *WrapReadWriteCloserConn) LocalAddr() net.Addr {
+	if conn.underConn != nil {
+		return conn.underConn.LocalAddr()
+	}
+	return (*net.TCPAddr)(nil)
+}
+
+func (conn *WrapReadWriteCloserConn) RemoteAddr() net.Addr {
+	if conn.underConn != nil {
+		return conn.underConn.RemoteAddr()
+	}
+	return (*net.TCPAddr)(nil)
+}
+
+func (conn *WrapReadWriteCloserConn) SetDeadline(t time.Time) error {
+	if conn.underConn != nil {
+		return conn.underConn.SetDeadline(t)
+	}
+	return &net.OpError{Op: "set", Net: "wrap", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
+}
+
+func (conn *WrapReadWriteCloserConn) SetReadDeadline(t time.Time) error {
+	if conn.underConn != nil {
+		return conn.underConn.SetReadDeadline(t)
+	}
+	return &net.OpError{Op: "set", Net: "wrap", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
+}
+
+func (conn *WrapReadWriteCloserConn) SetWriteDeadline(t time.Time) error {
+	if conn.underConn != nil {
+		return conn.underConn.SetWriteDeadline(t)
+	}
+	return &net.OpError{Op: "set", Net: "wrap", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
+}
+
+type CloseNotifyConn struct {
+	net.Conn
+
+	// 1 means closed
+	closeFlag int32
+
+	closeFn func()
+}
+
+// closeFn will be only called once
+func WrapCloseNotifyConn(c net.Conn, closeFn func()) net.Conn {
+	return &CloseNotifyConn{
+		Conn:    c,
+		closeFn: closeFn,
+	}
+}
+
+func (cc *CloseNotifyConn) Close() (err error) {
+	pflag := atomic.SwapInt32(&cc.closeFlag, 1)
+	if pflag == 0 {
+		err = cc.Close()
+		if cc.closeFn != nil {
+			cc.closeFn()
+		}
+	}
+	return
+}
+
+type StatsConn struct {
+	net.Conn
+
+	closed     int64 // 1 means closed
+	totalRead  int64
+	totalWrite int64
+	statsFunc  func(totalRead, totalWrite int64)
+}
+
+func WrapStatsConn(conn net.Conn, statsFunc func(total, totalWrite int64)) *StatsConn {
+	return &StatsConn{
+		Conn:      conn,
+		statsFunc: statsFunc,
+	}
+}
+
+func (statsConn *StatsConn) Read(p []byte) (n int, err error) {
+	n, err = statsConn.Conn.Read(p)
+	statsConn.totalRead += int64(n)
+	return
+}
+
+func (statsConn *StatsConn) Write(p []byte) (n int, err error) {
+	n, err = statsConn.Conn.Write(p)
+	statsConn.totalWrite += int64(n)
+	return
+}
+
+func (statsConn *StatsConn) Close() (err error) {
+	old := atomic.SwapInt64(&statsConn.closed, 1)
+	if old != 1 {
+		err = statsConn.Conn.Close()
+		if statsConn.statsFunc != nil {
+			statsConn.statsFunc(statsConn.totalRead, statsConn.totalWrite)
+		}
+	}
+	return
+}
+
+func ConnectServer(protocol string, addr string) (c net.Conn, err error) {
+	switch protocol {
+	case "tcp":
+		return net.Dial("tcp", addr)
+	case "kcp":
+		kcpConn, errRet := kcp.DialWithOptions(addr, nil, 10, 3)
+		if errRet != nil {
+			err = errRet
+			return
+		}
+		kcpConn.SetStreamMode(true)
+		kcpConn.SetWriteDelay(true)
+		kcpConn.SetNoDelay(1, 20, 2, 1)
+		kcpConn.SetWindowSize(128, 512)
+		kcpConn.SetMtu(1350)
+		kcpConn.SetACKNoDelay(false)
+		kcpConn.SetReadBuffer(4194304)
+		kcpConn.SetWriteBuffer(4194304)
+		c = kcpConn
+		return
+	default:
+		return nil, fmt.Errorf("unsupport protocol: %s", protocol)
+	}
+}
+
+func ConnectServerByProxy(proxyURL string, protocol string, addr string) (c net.Conn, err error) {
+	switch protocol {
+	case "tcp":
+		return gnet.DialTcpByProxy(proxyURL, addr)
+	case "kcp":
+		// http proxy is not supported for kcp
+		return ConnectServer(protocol, addr)
+	case "websocket":
+		return ConnectWebsocketServer(addr)
+	default:
+		return nil, fmt.Errorf("unsupport protocol: %s", protocol)
+	}
+}
+
+func ConnectServerByProxyWithTLS(proxyURL string, protocol string, addr string, tlsConfig *tls.Config) (c net.Conn, err error) {
+	c, err = ConnectServerByProxy(proxyURL, protocol, addr)
+	if err != nil {
+		return
+	}
+
+	if tlsConfig == nil {
+		return
+	}
+
+	c = WrapTLSClientConn(c, tlsConfig)
+	return
+}

pkg/util/net/http.go

@@ -0,0 +1,115 @@
+// Copyright 2017 fatedier, fatedier@gmail.com
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package net
+
+import (
+	"compress/gzip"
+	"io"
+	"net/http"
+	"strings"
+)
+
+type HTTPAuthWraper struct {
+	h      http.Handler
+	user   string
+	passwd string
+}
+
+func NewHTTPBasicAuthWraper(h http.Handler, user, passwd string) http.Handler {
+	return &HTTPAuthWraper{
+		h:      h,
+		user:   user,
+		passwd: passwd,
+	}
+}
+
+func (aw *HTTPAuthWraper) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	user, passwd, hasAuth := r.BasicAuth()
+	if (aw.user == "" && aw.passwd == "") || (hasAuth && user == aw.user && passwd == aw.passwd) {
+		aw.h.ServeHTTP(w, r)
+	} else {
+		w.Header().Set("WWW-Authenticate", `Basic realm="Restricted"`)
+		http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
+	}
+}
+
+type HTTPAuthMiddleware struct {
+	user   string
+	passwd string
+}
+
+func NewHTTPAuthMiddleware(user, passwd string) *HTTPAuthMiddleware {
+	return &HTTPAuthMiddleware{
+		user:   user,
+		passwd: passwd,
+	}
+}
+
+func (authMid *HTTPAuthMiddleware) Middleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		reqUser, reqPasswd, hasAuth := r.BasicAuth()
+		if (authMid.user == "" && authMid.passwd == "") ||
+			(hasAuth && reqUser == authMid.user && reqPasswd == authMid.passwd) {
+			next.ServeHTTP(w, r)
+		} else {
+			w.Header().Set("WWW-Authenticate", `Basic realm="Restricted"`)
+			http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
+		}
+	})
+}
+
+func HTTPBasicAuth(h http.HandlerFunc, user, passwd string) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		reqUser, reqPasswd, hasAuth := r.BasicAuth()
+		if (user == "" && passwd == "") ||
+			(hasAuth && reqUser == user && reqPasswd == passwd) {
+			h.ServeHTTP(w, r)
+		} else {
+			w.Header().Set("WWW-Authenticate", `Basic realm="Restricted"`)
+			http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
+		}
+	}
+}
+
+type HTTPGzipWraper struct {
+	h http.Handler
+}
+
+func (gw *HTTPGzipWraper) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
+		gw.h.ServeHTTP(w, r)
+		return
+	}
+	w.Header().Set("Content-Encoding", "gzip")
+	gz := gzip.NewWriter(w)
+	defer gz.Close()
+	gzr := gzipResponseWriter{Writer: gz, ResponseWriter: w}
+	gw.h.ServeHTTP(gzr, r)
+}
+
+func MakeHTTPGzipHandler(h http.Handler) http.Handler {
+	return &HTTPGzipWraper{
+		h: h,
+	}
+}
+
+type gzipResponseWriter struct {
+	io.Writer
+	http.ResponseWriter
+}
+
+func (w gzipResponseWriter) Write(b []byte) (int, error) {
+	return w.Writer.Write(b)
+}

pkg/util/net/kcp.go

@@ -0,0 +1,99 @@
+// Copyright 2017 fatedier, fatedier@gmail.com
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package net
+
+import (
+	"fmt"
+	"net"
+
+	kcp "github.com/fatedier/kcp-go"
+)
+
+type KCPListener struct {
+	listener  net.Listener
+	acceptCh  chan net.Conn
+	closeFlag bool
+}
+
+func ListenKcp(bindAddr string, bindPort int) (l *KCPListener, err error) {
+	listener, err := kcp.ListenWithOptions(fmt.Sprintf("%s:%d", bindAddr, bindPort), nil, 10, 3)
+	if err != nil {
+		return l, err
+	}
+	listener.SetReadBuffer(4194304)
+	listener.SetWriteBuffer(4194304)
+
+	l = &KCPListener{
+		listener:  listener,
+		acceptCh:  make(chan net.Conn),
+		closeFlag: false,
+	}
+
+	go func() {
+		for {
+			conn, err := listener.AcceptKCP()
+			if err != nil {
+				if l.closeFlag {
+					close(l.acceptCh)
+					return
+				}
+				continue
+			}
+			conn.SetStreamMode(true)
+			conn.SetWriteDelay(true)
+			conn.SetNoDelay(1, 20, 2, 1)
+			conn.SetMtu(1350)
+			conn.SetWindowSize(1024, 1024)
+			conn.SetACKNoDelay(false)
+
+			l.acceptCh <- conn
+		}
+	}()
+	return l, err
+}
+
+func (l *KCPListener) Accept() (net.Conn, error) {
+	conn, ok := <-l.acceptCh
+	if !ok {
+		return conn, fmt.Errorf("channel for kcp listener closed")
+	}
+	return conn, nil
+}
+
+func (l *KCPListener) Close() error {
+	if !l.closeFlag {
+		l.closeFlag = true
+		l.listener.Close()
+	}
+	return nil
+}
+
+func (l *KCPListener) Addr() net.Addr {
+	return l.listener.Addr()
+}
+
+func NewKCPConnFromUDP(conn *net.UDPConn, connected bool, raddr string) (net.Conn, error) {
+	kcpConn, err := kcp.NewConnEx(1, connected, raddr, nil, 10, 3, conn)
+	if err != nil {
+		return nil, err
+	}
+	kcpConn.SetStreamMode(true)
+	kcpConn.SetWriteDelay(true)
+	kcpConn.SetNoDelay(1, 20, 2, 1)
+	kcpConn.SetMtu(1350)
+	kcpConn.SetWindowSize(1024, 1024)
+	kcpConn.SetACKNoDelay(false)
+	return kcpConn, nil
+}

pkg/util/net/listener.go

@@ -0,0 +1,69 @@
+// Copyright 2017 fatedier, fatedier@gmail.com
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package net
+
+import (
+	"fmt"
+	"net"
+	"sync"
+
+	"github.com/fatedier/golib/errors"
+)
+
+// Custom listener
+type CustomListener struct {
+	acceptCh chan net.Conn
+	closed   bool
+	mu       sync.Mutex
+}
+
+func NewCustomListener() *CustomListener {
+	return &CustomListener{
+		acceptCh: make(chan net.Conn, 64),
+	}
+}
+
+func (l *CustomListener) Accept() (net.Conn, error) {
+	conn, ok := <-l.acceptCh
+	if !ok {
+		return nil, fmt.Errorf("listener closed")
+	}
+	return conn, nil
+}
+
+func (l *CustomListener) PutConn(conn net.Conn) error {
+	err := errors.PanicToError(func() {
+		select {
+		case l.acceptCh <- conn:
+		default:
+			conn.Close()
+		}
+	})
+	return err
+}
+
+func (l *CustomListener) Close() error {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	if !l.closed {
+		close(l.acceptCh)
+		l.closed = true
+	}
+	return nil
+}
+
+func (l *CustomListener) Addr() net.Addr {
+	return (*net.TCPAddr)(nil)
+}

pkg/util/net/tls.go

@@ -0,0 +1,57 @@
+// Copyright 2019 fatedier, fatedier@gmail.com
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package net
+
+import (
+	"crypto/tls"
+	"fmt"
+	"net"
+	"time"
+
+	gnet "github.com/fatedier/golib/net"
+)
+
+var (
+	FRPTLSHeadByte = 0x17
+)
+
+func WrapTLSClientConn(c net.Conn, tlsConfig *tls.Config) (out net.Conn) {
+	c.Write([]byte{byte(FRPTLSHeadByte)})
+	out = tls.Client(c, tlsConfig)
+	return
+}
+
+func CheckAndEnableTLSServerConnWithTimeout(c net.Conn, tlsConfig *tls.Config, tlsOnly bool, timeout time.Duration) (out net.Conn, err error) {
+	sc, r := gnet.NewSharedConnSize(c, 2)
+	buf := make([]byte, 1)
+	var n int
+	c.SetReadDeadline(time.Now().Add(timeout))
+	n, err = r.Read(buf)
+	c.SetReadDeadline(time.Time{})
+	if err != nil {
+		return
+	}
+
+	if n == 1 && int(buf[0]) == FRPTLSHeadByte {
+		out = tls.Server(c, tlsConfig)
+	} else {
+		if tlsOnly {
+			err = fmt.Errorf("non-TLS connection received on a TlsOnly server")
+			return
+		}
+		out = sc
+	}
+	return
+}

pkg/util/net/udp.go

@@ -0,0 +1,258 @@
+// Copyright 2017 fatedier, fatedier@gmail.com
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package net
+
+import (
+	"fmt"
+	"io"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/fatedier/golib/pool"
+)
+
+type UDPPacket struct {
+	Buf        []byte
+	LocalAddr  net.Addr
+	RemoteAddr net.Addr
+}
+
+type FakeUDPConn struct {
+	l *UDPListener
+
+	localAddr  net.Addr
+	remoteAddr net.Addr
+	packets    chan []byte
+	closeFlag  bool
+
+	lastActive time.Time
+	mu         sync.RWMutex
+}
+
+func NewFakeUDPConn(l *UDPListener, laddr, raddr net.Addr) *FakeUDPConn {
+	fc := &FakeUDPConn{
+		l:          l,
+		localAddr:  laddr,
+		remoteAddr: raddr,
+		packets:    make(chan []byte, 20),
+	}
+
+	go func() {
+		for {
+			time.Sleep(5 * time.Second)
+			fc.mu.RLock()
+			if time.Now().Sub(fc.lastActive) > 10*time.Second {
+				fc.mu.RUnlock()
+				fc.Close()
+				break
+			}
+			fc.mu.RUnlock()
+		}
+	}()
+	return fc
+}
+
+func (c *FakeUDPConn) putPacket(content []byte) {
+	defer func() {
+		if err := recover(); err != nil {
+		}
+	}()
+
+	select {
+	case c.packets <- content:
+	default:
+	}
+}
+
+func (c *FakeUDPConn) Read(b []byte) (n int, err error) {
+	content, ok := <-c.packets
+	if !ok {
+		return 0, io.EOF
+	}
+	c.mu.Lock()
+	c.lastActive = time.Now()
+	c.mu.Unlock()
+
+	if len(b) < len(content) {
+		n = len(b)
+	} else {
+		n = len(content)
+	}
+	copy(b, content)
+	return n, nil
+}
+
+func (c *FakeUDPConn) Write(b []byte) (n int, err error) {
+	c.mu.RLock()
+	if c.closeFlag {
+		c.mu.RUnlock()
+		return 0, io.ErrClosedPipe
+	}
+	c.mu.RUnlock()
+
+	packet := &UDPPacket{
+		Buf:        b,
+		LocalAddr:  c.localAddr,
+		RemoteAddr: c.remoteAddr,
+	}
+	c.l.writeUDPPacket(packet)
+
+	c.mu.Lock()
+	c.lastActive = time.Now()
+	c.mu.Unlock()
+	return len(b), nil
+}
+
+func (c *FakeUDPConn) Close() error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if !c.closeFlag {
+		c.closeFlag = true
+		close(c.packets)
+	}
+	return nil
+}
+
+func (c *FakeUDPConn) IsClosed() bool {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.closeFlag
+}
+
+func (c *FakeUDPConn) LocalAddr() net.Addr {
+	return c.localAddr
+}
+
+func (c *FakeUDPConn) RemoteAddr() net.Addr {
+	return c.remoteAddr
+}
+
+func (c *FakeUDPConn) SetDeadline(t time.Time) error {
+	return nil
+}
+
+func (c *FakeUDPConn) SetReadDeadline(t time.Time) error {
+	return nil
+}
+
+func (c *FakeUDPConn) SetWriteDeadline(t time.Time) error {
+	return nil
+}
+
+type UDPListener struct {
+	addr      net.Addr
+	acceptCh  chan net.Conn
+	writeCh   chan *UDPPacket
+	readConn  net.Conn
+	closeFlag bool
+
+	fakeConns map[string]*FakeUDPConn
+}
+
+func ListenUDP(bindAddr string, bindPort int) (l *UDPListener, err error) {
+	udpAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", bindAddr, bindPort))
+	if err != nil {
+		return l, err
+	}
+	readConn, err := net.ListenUDP("udp", udpAddr)
+
+	l = &UDPListener{
+		addr:      udpAddr,
+		acceptCh:  make(chan net.Conn),
+		writeCh:   make(chan *UDPPacket, 1000),
+		fakeConns: make(map[string]*FakeUDPConn),
+	}
+
+	// for reading
+	go func() {
+		for {
+			buf := pool.GetBuf(1450)
+			n, remoteAddr, err := readConn.ReadFromUDP(buf)
+			if err != nil {
+				close(l.acceptCh)
+				close(l.writeCh)
+				return
+			}
+
+			fakeConn, exist := l.fakeConns[remoteAddr.String()]
+			if !exist || fakeConn.IsClosed() {
+				fakeConn = NewFakeUDPConn(l, l.Addr(), remoteAddr)
+				l.fakeConns[remoteAddr.String()] = fakeConn
+			}
+			fakeConn.putPacket(buf[:n])
+
+			l.acceptCh <- fakeConn
+		}
+	}()
+
+	// for writing
+	go func() {
+		for {
+			packet, ok := <-l.writeCh
+			if !ok {
+				return
+			}
+
+			if addr, ok := packet.RemoteAddr.(*net.UDPAddr); ok {
+				readConn.WriteToUDP(packet.Buf, addr)
+			}
+		}
+	}()
+
+	return
+}
+
+func (l *UDPListener) writeUDPPacket(packet *UDPPacket) (err error) {
+	defer func() {
+		if errRet := recover(); errRet != nil {
+			err = fmt.Errorf("udp write closed listener")
+		}
+	}()
+	l.writeCh <- packet
+	return
+}
+
+func (l *UDPListener) WriteMsg(buf []byte, remoteAddr *net.UDPAddr) (err error) {
+	// only set remote addr here
+	packet := &UDPPacket{
+		Buf:        buf,
+		RemoteAddr: remoteAddr,
+	}
+	err = l.writeUDPPacket(packet)
+	return
+}
+
+func (l *UDPListener) Accept() (net.Conn, error) {
+	conn, ok := <-l.acceptCh
+	if !ok {
+		return conn, fmt.Errorf("channel for udp listener closed")
+	}
+	return conn, nil
+}
+
+func (l *UDPListener) Close() error {
+	if !l.closeFlag {
+		l.closeFlag = true
+		if l.readConn != nil {
+			l.readConn.Close()
+		}
+	}
+	return nil
+}
+
+func (l *UDPListener) Addr() net.Addr {
+	return l.addr
+}

pkg/util/net/websocket.go

@@ -0,0 +1,103 @@
+package net
+
+import (
+	"errors"
+	"fmt"
+	"net"
+	"net/http"
+	"net/url"
+	"time"
+
+	"golang.org/x/net/websocket"
+)
+
+var (
+	ErrWebsocketListenerClosed = errors.New("websocket listener closed")
+)
+
+const (
+	FrpWebsocketPath = "/~!frp"
+)
+
+type WebsocketListener struct {
+	ln       net.Listener
+	acceptCh chan net.Conn
+
+	server    *http.Server
+	httpMutex *http.ServeMux
+}
+
+// NewWebsocketListener to handle websocket connections
+// ln: tcp listener for websocket connections
+func NewWebsocketListener(ln net.Listener) (wl *WebsocketListener) {
+	wl = &WebsocketListener{
+		acceptCh: make(chan net.Conn),
+	}
+
+	muxer := http.NewServeMux()
+	muxer.Handle(FrpWebsocketPath, websocket.Handler(func(c *websocket.Conn) {
+		notifyCh := make(chan struct{})
+		conn := WrapCloseNotifyConn(c, func() {
+			close(notifyCh)
+		})
+		wl.acceptCh <- conn
+		<-notifyCh
+	}))
+
+	wl.server = &http.Server{
+		Addr:    ln.Addr().String(),
+		Handler: muxer,
+	}
+
+	go wl.server.Serve(ln)
+	return
+}
+
+func ListenWebsocket(bindAddr string, bindPort int) (*WebsocketListener, error) {
+	tcpLn, err := net.Listen("tcp", fmt.Sprintf("%s:%d", bindAddr, bindPort))
+	if err != nil {
+		return nil, err
+	}
+	l := NewWebsocketListener(tcpLn)
+	return l, nil
+}
+
+func (p *WebsocketListener) Accept() (net.Conn, error) {
+	c, ok := <-p.acceptCh
+	if !ok {
+		return nil, ErrWebsocketListenerClosed
+	}
+	return c, nil
+}
+
+func (p *WebsocketListener) Close() error {
+	return p.server.Close()
+}
+
+func (p *WebsocketListener) Addr() net.Addr {
+	return p.ln.Addr()
+}
+
+// addr: domain:port
+func ConnectWebsocketServer(addr string) (net.Conn, error) {
+	addr = "ws://" + addr + FrpWebsocketPath
+	uri, err := url.Parse(addr)
+	if err != nil {
+		return nil, err
+	}
+
+	origin := "http://" + 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
+	}
+	return conn, nil
+}