v2fly/transport/internet/kcp/connection.go

package kcp

import (
	"errors"
	"io"
	"net"
	"sync"
	"time"

	"github.com/v2ray/v2ray-core/common/alloc"
	"github.com/v2ray/v2ray-core/common/log"
	"github.com/v2ray/v2ray-core/common/signal"
)

var (
	errTimeout          = errors.New("i/o timeout")
	errBrokenPipe       = errors.New("broken pipe")
	errClosedListener   = errors.New("Listener closed.")
	errClosedConnection = errors.New("Connection closed.")
)

const (
	headerSize uint32 = 2
)

type Command byte

var (
	CommandData      Command = 0
	CommandTerminate Command = 1
)

type Option byte

var (
	OptionClose Option = 1
)

type ConnState byte

var (
	ConnStateActive       ConnState = 0
	ConnStateReadyToClose ConnState = 1
	ConnStatePeerClosed   ConnState = 2
	ConnStateClosed       ConnState = 4
)

func nowMillisec() int64 {
	now := time.Now()
	return now.Unix()*1000 + int64(now.Nanosecond()/1000000)
}

// Connection is a KCP connection over UDP.
type Connection struct {
	sync.RWMutex
	state           ConnState
	kcp             *KCP // the core ARQ
	kcpAccess       sync.Mutex
	block           Authenticator
	needUpdate      bool
	local, remote   net.Addr
	rd              time.Time // read deadline
	wd              time.Time // write deadline
	chReadEvent     chan struct{}
	writer          io.WriteCloser
	since           int64
	terminateOnce   signal.Once
	writeBufferSize uint32
}

// NewConnection create a new KCP connection between local and remote.
func NewConnection(conv uint32, writerCloser io.WriteCloser, local *net.UDPAddr, remote *net.UDPAddr, block Authenticator) *Connection {
	conn := new(Connection)
	conn.local = local
	conn.chReadEvent = make(chan struct{}, 1)
	conn.remote = remote
	conn.block = block
	conn.writer = writerCloser
	conn.since = nowMillisec()

	mtu := effectiveConfig.Mtu - uint32(block.HeaderSize()) - headerSize
	conn.kcp = NewKCP(conv, mtu, effectiveConfig.GetSendingWindowSize(), effectiveConfig.GetReceivingWindowSize(), conn.output)
	conn.kcp.NoDelay(effectiveConfig.Tti, 2, effectiveConfig.Congestion)
	conn.kcp.current = conn.Elapsed()
	conn.writeBufferSize = effectiveConfig.WriteBuffer / effectiveConfig.Mtu

	go conn.updateTask()

	return conn
}

func (this *Connection) Elapsed() uint32 {
	return uint32(nowMillisec() - this.since)
}

// Read implements the Conn Read method.
func (this *Connection) Read(b []byte) (int, error) {
	if this == nil || this.state == ConnStateReadyToClose || this.state == ConnStateClosed {
		return 0, io.EOF
	}

	for {
		this.RLock()
		if this.state == ConnStateReadyToClose || this.state == ConnStateClosed {
			this.RUnlock()
			return 0, io.EOF
		}

		if !this.rd.IsZero() && this.rd.Before(time.Now()) {
			this.RUnlock()
			return 0, errTimeout
		}
		this.RUnlock()

		this.kcpAccess.Lock()
		nBytes := this.kcp.Recv(b)
		this.kcpAccess.Unlock()
		if nBytes > 0 {
			return nBytes, nil
		}
		select {
		case <-this.chReadEvent:
		case <-time.After(time.Second):
		}
	}
}

// Write implements the Conn Write method.
func (this *Connection) Write(b []byte) (int, error) {
	if this == nil ||
		this.state == ConnStateReadyToClose ||
		this.state == ConnStatePeerClosed ||
		this.state == ConnStateClosed {
		return 0, io.ErrClosedPipe
	}

	for {
		this.RLock()
		if this.state == ConnStateReadyToClose ||
			this.state == ConnStatePeerClosed ||
			this.state == ConnStateClosed {
			this.RUnlock()
			return 0, io.ErrClosedPipe
		}
		this.RUnlock()

		this.kcpAccess.Lock()
		if this.kcp.WaitSnd() < this.writeBufferSize {
			nBytes := len(b)
			this.kcp.Send(b)
			this.kcp.current = this.Elapsed()
			this.kcp.flush()
			this.kcpAccess.Unlock()
			return nBytes, nil
		}
		this.kcpAccess.Unlock()

		if !this.wd.IsZero() && this.wd.Before(time.Now()) {
			return 0, errTimeout
		}

		// Sending windows is 1024 for the moment. This amount is not gonna sent in 1 sec.
		time.Sleep(time.Second)
	}
}

func (this *Connection) Terminate() {
	if this == nil || this.state == ConnStateClosed {
		return
	}
	this.Lock()
	defer this.Unlock()
	if this.state == ConnStateClosed {
		return
	}

	this.state = ConnStateClosed
	this.writer.Close()
}

func (this *Connection) NotifyTermination() {
	for i := 0; i < 16; i++ {
		this.RLock()
		if this.state == ConnStateClosed {
			this.RUnlock()
			break
		}
		this.RUnlock()
		buffer := alloc.NewSmallBuffer().Clear()
		buffer.AppendBytes(byte(CommandTerminate), byte(OptionClose), byte(0), byte(0), byte(0), byte(0))
		this.outputBuffer(buffer)

		time.Sleep(time.Second)

	}
	this.Terminate()
}

func (this *Connection) ForceTimeout() {
	if this == nil {
		return
	}
	for i := 0; i < 5; i++ {
		if this.state == ConnStateClosed {
			return
		}
		time.Sleep(time.Minute)
	}
	go this.terminateOnce.Do(this.NotifyTermination)
}

// Close closes the connection.
func (this *Connection) Close() error {
	if this == nil || this.state == ConnStateClosed || this.state == ConnStateReadyToClose {
		return errClosedConnection
	}
	log.Debug("KCP|Connection: Closing connection to ", this.remote)
	this.Lock()
	defer this.Unlock()

	if this.state == ConnStateActive {
		this.state = ConnStateReadyToClose
		if this.kcp.WaitSnd() == 0 {
			go this.terminateOnce.Do(this.NotifyTermination)
		} else {
			go this.ForceTimeout()
		}
	}

	if this.state == ConnStatePeerClosed {
		go this.Terminate()
	}

	return nil
}

// LocalAddr returns the local network address. The Addr returned is shared by all invocations of LocalAddr, so do not modify it.
func (this *Connection) LocalAddr() net.Addr {
	if this == nil {
		return nil
	}
	return this.local
}

// RemoteAddr returns the remote network address. The Addr returned is shared by all invocations of RemoteAddr, so do not modify it.
func (this *Connection) RemoteAddr() net.Addr {
	if this == nil {
		return nil
	}
	return this.remote
}

// SetDeadline sets the deadline associated with the listener. A zero time value disables the deadline.
func (this *Connection) SetDeadline(t time.Time) error {
	if this == nil || this.state != ConnStateActive {
		return errClosedConnection
	}
	this.Lock()
	defer this.Unlock()
	this.rd = t
	this.wd = t
	return nil
}

// SetReadDeadline implements the Conn SetReadDeadline method.
func (this *Connection) SetReadDeadline(t time.Time) error {
	if this == nil || this.state != ConnStateActive {
		return errClosedConnection
	}
	this.Lock()
	defer this.Unlock()
	this.rd = t
	return nil
}

// SetWriteDeadline implements the Conn SetWriteDeadline method.
func (this *Connection) SetWriteDeadline(t time.Time) error {
	if this == nil || this.state != ConnStateActive {
		return errClosedConnection
	}
	this.Lock()
	defer this.Unlock()
	this.wd = t
	return nil
}

func (this *Connection) outputBuffer(payload *alloc.Buffer) {
	defer payload.Release()
	if this == nil {
		return
	}

	this.RLock()
	defer this.RUnlock()
	if this.state == ConnStatePeerClosed || this.state == ConnStateClosed {
		return
	}
	this.block.Seal(payload)

	this.writer.Write(payload.Value)
}

func (this *Connection) output(payload []byte) {
	if this == nil || this.state == ConnStateClosed {
		return
	}

	if this.state == ConnStateReadyToClose && this.kcp.WaitSnd() == 0 {
		go this.terminateOnce.Do(this.NotifyTermination)
	}

	if len(payload) < IKCP_OVERHEAD {
		return
	}

	buffer := alloc.NewBuffer().Clear().Append(payload)
	cmd := CommandData
	opt := Option(0)
	if this.state == ConnStateReadyToClose {
		opt = OptionClose
	}
	buffer.Prepend([]byte{byte(cmd), byte(opt)})
	this.outputBuffer(buffer)
}

// kcp update, input loop
func (this *Connection) updateTask() {
	for this.state != ConnStateClosed {
		current := this.Elapsed()
		this.kcpAccess.Lock()
		this.kcp.Update(current)
		this.kcpAccess.Unlock()
		time.Sleep(time.Duration(effectiveConfig.Tti) * time.Millisecond)
	}
}

func (this *Connection) notifyReadEvent() {
	select {
	case this.chReadEvent <- struct{}{}:
	default:
	}
}

func (this *Connection) MarkPeerClose() {
	this.Lock()
	defer this.Unlock()
	if this.state == ConnStateReadyToClose {
		this.state = ConnStateClosed
		go this.Terminate()
		return
	}
	if this.state == ConnStateActive {
		this.state = ConnStatePeerClosed
	}
	this.kcpAccess.Lock()
	this.kcp.ClearSendQueue()
	this.kcpAccess.Unlock()
}

func (this *Connection) kcpInput(data []byte) {
	cmd := Command(data[0])
	opt := Option(data[1])
	if cmd == CommandTerminate {
		go this.Terminate()
		return
	}
	if opt == OptionClose {
		go this.MarkPeerClose()
	}
	this.kcpAccess.Lock()
	this.kcp.current = this.Elapsed()
	this.kcp.Input(data[2:])

	this.kcpAccess.Unlock()
	this.notifyReadEvent()
}

func (this *Connection) FetchInputFrom(conn net.Conn) {
	go func() {
		for {
			payload := alloc.NewBuffer()
			nBytes, err := conn.Read(payload.Value)
			if err != nil {
				return
			}
			payload.Slice(0, nBytes)
			if this.block.Open(payload) {
				this.kcpInput(payload.Value)
			} else {
				log.Info("KCP|Connection: Invalid response from ", conn.RemoteAddr())
			}
			payload.Release()
		}
	}()
}

func (this *Connection) Reusable() bool {
	return false
}

func (this *Connection) SetReusable(b bool) {}