package core import ( "context" "crypto/tls" "errors" "fmt" "github.com/lucas-clemente/quic-go" "github.com/lucas-clemente/quic-go/congestion" "github.com/lunixbochs/struc" "net" "sync" "time" ) var ( ErrClosed = errors.New("client closed") ) type CongestionFactory func(refBPS uint64) congestion.CongestionControl type Client struct { serverAddr string sendBPS, recvBPS uint64 auth []byte congestionFactory CongestionFactory obfuscator Obfuscator tlsConfig *tls.Config quicConfig *quic.Config quicSession quic.Session reconnectMutex sync.Mutex closed bool } func NewClient(serverAddr string, auth []byte, tlsConfig *tls.Config, quicConfig *quic.Config, sendBPS uint64, recvBPS uint64, congestionFactory CongestionFactory, obfuscator Obfuscator) (*Client, error) { c := &Client{ serverAddr: serverAddr, sendBPS: sendBPS, recvBPS: recvBPS, auth: auth, congestionFactory: congestionFactory, obfuscator: obfuscator, tlsConfig: tlsConfig, quicConfig: quicConfig, } if err := c.connectToServer(); err != nil { return nil, err } return c, nil } func (c *Client) connectToServer() error { serverUDPAddr, err := net.ResolveUDPAddr("udp", c.serverAddr) if err != nil { return err } packetConn, err := net.ListenPacket("udp", "") if err != nil { return err } if c.obfuscator != nil { // Wrap PacketConn with obfuscator packetConn = &obfsPacketConn{ Orig: packetConn, Obfuscator: c.obfuscator, } } qs, err := quic.Dial(packetConn, serverUDPAddr, c.serverAddr, c.tlsConfig, c.quicConfig) if err != nil { return err } // Control stream ctx, ctxCancel := context.WithTimeout(context.Background(), protocolTimeout) stream, err := qs.OpenStreamSync(ctx) ctxCancel() if err != nil { _ = qs.CloseWithError(closeErrorCodeProtocol, "protocol error") return err } ok, msg, err := c.handleControlStream(qs, stream) if err != nil { _ = qs.CloseWithError(closeErrorCodeProtocol, "protocol error") return err } if !ok { _ = qs.CloseWithError(closeErrorCodeAuth, "auth error") return fmt.Errorf("auth error: %s", msg) } // All good c.quicSession = qs return nil } func (c *Client) handleControlStream(qs quic.Session, stream quic.Stream) (bool, string, error) { // Send client hello err := struc.Pack(stream, &clientHello{ Rate: transmissionRate{ SendBPS: c.sendBPS, RecvBPS: c.recvBPS, }, Auth: c.auth, }) if err != nil { return false, "", err } // Receive server hello var sh serverHello err = struc.Unpack(stream, &sh) if err != nil { return false, "", err } // Set the congestion accordingly if sh.OK && c.congestionFactory != nil { qs.SetCongestionControl(c.congestionFactory(sh.Rate.RecvBPS)) } return true, sh.Message, nil } func (c *Client) openStreamWithReconnect() (quic.Stream, net.Addr, net.Addr, error) { c.reconnectMutex.Lock() defer c.reconnectMutex.Unlock() if c.closed { return nil, nil, nil, ErrClosed } stream, err := c.quicSession.OpenStream() if err == nil { // All good return stream, c.quicSession.LocalAddr(), c.quicSession.RemoteAddr(), nil } // Something is wrong if nErr, ok := err.(net.Error); ok && nErr.Temporary() { // Temporary error, just return return nil, nil, nil, err } // Permanent error, need to reconnect if err := c.connectToServer(); err != nil { // Still error, oops return nil, nil, nil, err } // We are not going to try again even if it still fails the second time stream, err = c.quicSession.OpenStream() return stream, c.quicSession.LocalAddr(), c.quicSession.RemoteAddr(), err } func (c *Client) DialTCP(addr string) (net.Conn, error) { stream, localAddr, remoteAddr, err := c.openStreamWithReconnect() if err != nil { return nil, err } // Send request err = struc.Pack(stream, &clientRequest{ UDP: false, Address: addr, }) if err != nil { _ = stream.Close() return nil, err } // Read response var sr serverResponse err = struc.Unpack(stream, &sr) if err != nil { _ = stream.Close() return nil, err } if !sr.OK { _ = stream.Close() return nil, fmt.Errorf("connection rejected: %s", sr.Message) } return &quicConn{ Orig: stream, PseudoLocalAddr: localAddr, PseudoRemoteAddr: remoteAddr, }, nil } func (c *Client) Close() error { c.reconnectMutex.Lock() defer c.reconnectMutex.Unlock() err := c.quicSession.CloseWithError(closeErrorCodeGeneric, "") c.closed = true return err } type quicConn struct { Orig quic.Stream PseudoLocalAddr net.Addr PseudoRemoteAddr net.Addr } func (w *quicConn) Read(b []byte) (n int, err error) { return w.Orig.Read(b) } func (w *quicConn) Write(b []byte) (n int, err error) { return w.Orig.Write(b) } func (w *quicConn) Close() error { return w.Orig.Close() } func (w *quicConn) LocalAddr() net.Addr { return w.PseudoLocalAddr } func (w *quicConn) RemoteAddr() net.Addr { return w.PseudoRemoteAddr } func (w *quicConn) SetDeadline(t time.Time) error { return w.Orig.SetDeadline(t) } func (w *quicConn) SetReadDeadline(t time.Time) error { return w.Orig.SetReadDeadline(t) } func (w *quicConn) SetWriteDeadline(t time.Time) error { return w.Orig.SetWriteDeadline(t) }