Xray-core/transport/internet/finalmask/xmc/deadline.go

package xmc

import (
	"net"
	"sync"
	"time"
)

const handshakeTimeout = 2 * time.Minute

type connectionDeadlines struct {
	mu sync.Mutex
	c  net.Conn

	read      time.Time
	write     time.Time
	handshake time.Time
}

func newConnectionDeadlines(c net.Conn) *connectionDeadlines {
	return &connectionDeadlines{c: c}
}

func (d *connectionDeadlines) beginHandshake() error {
	d.mu.Lock()
	defer d.mu.Unlock()

	d.handshake = time.Now().Add(handshakeTimeout)
	if err := d.applyLocked(); err != nil {
		d.handshake = time.Time{}
		_ = d.applyLocked()
		return err
	}
	return nil
}

func (d *connectionDeadlines) endHandshake() error {
	d.mu.Lock()
	defer d.mu.Unlock()

	d.handshake = time.Time{}
	return d.applyLocked()
}

func (d *connectionDeadlines) setDeadline(t time.Time) error {
	d.mu.Lock()
	defer d.mu.Unlock()

	d.read = t
	d.write = t
	return d.applyLocked()
}

func (d *connectionDeadlines) setReadDeadline(t time.Time) error {
	d.mu.Lock()
	defer d.mu.Unlock()

	d.read = t
	return d.c.SetReadDeadline(earlierDeadline(d.read, d.handshake))
}

func (d *connectionDeadlines) setWriteDeadline(t time.Time) error {
	d.mu.Lock()
	defer d.mu.Unlock()

	d.write = t
	return d.c.SetWriteDeadline(earlierDeadline(d.write, d.handshake))
}

func (d *connectionDeadlines) applyLocked() error {
	if err := d.c.SetReadDeadline(earlierDeadline(d.read, d.handshake)); err != nil {
		return err
	}
	return d.c.SetWriteDeadline(earlierDeadline(d.write, d.handshake))
}

func earlierDeadline(user, internal time.Time) time.Time {
	if internal.IsZero() {
		return user
	}
	if user.IsZero() || internal.Before(user) {
		return internal
	}
	return user
}