Xray-core/transport/internet/websocket/connection.go

package websocket

import (
	"context"
	"io"
	"net"
	"time"

	"github.com/coder/websocket"
	"github.com/xtls/xray-core/common/buf"
)

var _ buf.Writer = (*connection)(nil)

type connection struct {
	net.Conn
	reader     io.Reader
	remoteAddr net.Addr
	localAddr  net.Addr
}

func NewConnection(conn *websocket.Conn, remoteAddr net.Addr, extraReader io.Reader, heartbeatPeriod uint32) *connection {
	ctx, cancel := context.WithCancel(context.Background())
	c := &connection{Conn: websocket.NetConn(ctx, conn, websocket.MessageBinary), reader: extraReader, remoteAddr: remoteAddr}
	if remoteAddr == nil {
		c.remoteAddr = c.Conn.RemoteAddr()
	}
	c.Conn = &cancelConn{Conn: c.Conn, cancel: cancel}
	if heartbeatPeriod != 0 {
		go func() {
			ticker := time.NewTicker(time.Duration(heartbeatPeriod) * time.Second)
			defer ticker.Stop()
			for {
				select {
				case <-ctx.Done():
					return
				case <-ticker.C:
					pingCtx, stop := context.WithTimeout(ctx, time.Duration(heartbeatPeriod)*time.Second)
					err := conn.Ping(pingCtx)
					stop()
					if err != nil {
						c.Close()
						return
					}
				}
			}
		}()
	}
	return c
}

type cancelConn struct {
	net.Conn
	cancel context.CancelFunc
}

func (c *cancelConn) Close() error {
	defer c.cancel()
	return c.Conn.Close()
}

func (c *connection) Read(b []byte) (int, error) {
	if c.reader != nil {
		n, err := c.reader.Read(b)
		if err == io.EOF {
			c.reader = nil
			if n > 0 {
				return n, nil
			}
		} else {
			return n, err
		}
	}
	return c.Conn.Read(b)
}

func (c *connection) LocalAddr() net.Addr {
	if c.localAddr != nil {
		return c.localAddr
	}
	return c.Conn.LocalAddr()
}

func (c *connection) RemoteAddr() net.Addr { return c.remoteAddr }

func (c *connection) WriteMultiBuffer(mb buf.MultiBuffer) error {
	mb = buf.Compact(mb)
	mb, err := buf.WriteMultiBuffer(c, mb)
	buf.ReleaseMulti(mb)
	return err
}