connection.go

package connection

import (
	"fmt"
	"net"
	"strconv"
	"sync"
	"time"

	"go.aporeto.io/trireme-lib/controller/internal/enforcer/nfqdatapath/afinetrawsocket"
	"go.aporeto.io/trireme-lib/controller/pkg/claimsheader"
	"go.aporeto.io/trireme-lib/controller/pkg/counters"
	"go.aporeto.io/trireme-lib/controller/pkg/packet"
	"go.aporeto.io/trireme-lib/controller/pkg/pucontext"
	"go.aporeto.io/trireme-lib/controller/pkg/secrets"
	"go.aporeto.io/trireme-lib/policy"
	"go.aporeto.io/trireme-lib/utils/cache"
	"go.aporeto.io/trireme-lib/utils/crypto"
	"go.uber.org/zap"
)

// TCPFlowState identifies the constants of the state of a TCP connectioncon
type TCPFlowState int

// UDPFlowState identifies the constants of the state of a UDP connection.
type UDPFlowState int

// ProxyConnState identifies the constants of the state of a proxied connection
type ProxyConnState int

const (

	// TCPSynSend is the state where the Syn packets has been send, but no response has been received
	TCPSynSend TCPFlowState = iota

	// TCPSynReceived indicates that the syn packet has been received
	TCPSynReceived

	// TCPSynAckSend indicates that the SynAck packet has been send
	TCPSynAckSend

	// TCPSynAckReceived is the state where the SynAck has been received
	TCPSynAckReceived

	// TCPAckSend indicates that the ack packets has been sent
	TCPAckSend

	// TCPAckProcessed is the state that the negotiation has been completed
	TCPAckProcessed

	// TCPData indicates that the packets are now data packets
	TCPData

	// UnknownState indicates that this an existing connection in the uknown state.
	UnknownState
)

const (
	// ClientTokenSend Init token send for client
	ClientTokenSend ProxyConnState = iota

	// ServerReceivePeerToken -- waiting to receive peer token
	ServerReceivePeerToken

	// ServerSendToken -- Send our own token and the client tokens
	ServerSendToken

	// ClientPeerTokenReceive -- Receive signed tokens from server
	ClientPeerTokenReceive

	// ClientSendSignedPair -- Sign the (token/nonce pair) and send
	ClientSendSignedPair

	// ServerAuthenticatePair -- Authenticate pair of tokens
	ServerAuthenticatePair
)

const (
	// UDPStart is the state where a syn will be sent.
	UDPStart UDPFlowState = iota

	// UDPClientSendSyn is the state where a syn has been sent.
	UDPClientSendSyn

	// UDPClientSendAck  is the state where application side has send the ACK.
	UDPClientSendAck

	// UDPReceiverSendSynAck is the state where syn ack packet has been sent.
	UDPReceiverSendSynAck

	// UDPReceiverProcessedAck is the state that the negotiation has been completed.
	UDPReceiverProcessedAck

	// UDPData is the state where data is being transmitted.
	UDPData
)

// MaximumUDPQueueLen is the maximum number of UDP packets buffered.
const MaximumUDPQueueLen = 50

// AuthInfo keeps authentication information about a connection
type AuthInfo struct {
	LocalContext         []byte
	RemoteContext        []byte
	RemoteContextID      string
	RemotePublicKey      interface{}
	RemoteIP             string
	RemotePort           string
	LocalServiceContext  []byte
	RemoteServiceContext []byte
}

//TCPTuple contains the 4 tuple for tcp connection
type TCPTuple struct {
	SourceAddress      net.IP
	DestinationAddress net.IP
	SourcePort         uint16
	DestinationPort    uint16
}

// PingConfig holds ping configuration per connection.
type PingConfig struct {
	StartTime   time.Time
	Type        claimsheader.PingType
	Passthrough bool
	SessionID   string
	Request     int
}

// TCPConnection is information regarding TCP Connection
type TCPConnection struct {
	sync.RWMutex

	state TCPFlowState
	Auth  AuthInfo

	// ServiceData allows services to associate state with a connection
	ServiceData interface{}

	// Context is the pucontext.PUContext that is associated with this connection
	// Minimizes the number of caches and lookups
	Context *pucontext.PUContext

	// TimeOut signals the timeout to be used by the state machines
	TimeOut time.Duration

	// ServiceConnection indicates that this connection is handled by a service
	ServiceConnection bool

	// LoopbackConnection indicates that this connections is within the same pu context.
	loopbackConnection bool

	// ReportFlowPolicy holds the last matched observed policy
	ReportFlowPolicy *policy.FlowPolicy

	// PacketFlowPolicy holds the last matched actual policy
	PacketFlowPolicy *policy.FlowPolicy

	// MarkForDeletion -- this is is used only in conjunction with serviceconnection. Its a hint for us if we have a fin for an earlier connection
	// and this is reused port flow.
	MarkForDeletion bool

	RetransmittedSynAck bool

	expiredConnection bool

	// TCPtuple is tcp tuple
	TCPtuple *TCPTuple

	// PingConfig is the config of the enforcer ping diagnostic feature
	PingConfig *PingConfig

	Secrets secrets.Secrets

	SourceController      string
	DestinationController string
	initialSequenceNumber uint32
}

// TCPConnectionExpirationNotifier handles processing the expiration of an element
func TCPConnectionExpirationNotifier(_ interface{}, item interface{}) {

	conn, ok := item.(*TCPConnection)
	if !ok {
		return
	}

	conn.Cleanup()
}

func (tcpTuple *TCPTuple) String() string {
	return "sip: " + tcpTuple.SourceAddress.String() + " dip: " + tcpTuple.DestinationAddress.String() + " sport: " + strconv.Itoa(int(tcpTuple.SourcePort)) + " dport: " + strconv.Itoa(int(tcpTuple.DestinationPort))
}

// String returns a printable version of connection
func (c *TCPConnection) String() string {
	return fmt.Sprintf("state:%d auth: %+v", c.state, c.Auth)
}

// GetState is used to return the state
func (c *TCPConnection) GetState() TCPFlowState {
	return c.state
}

// GetInitialSequenceNumber returns the initial sequence number that was found on the syn packet
// corresponding to this TCP Connection
func (c *TCPConnection) GetInitialSequenceNumber() uint32 {
	return c.initialSequenceNumber
}

// GetMarkForDeletion returns the state of markForDeletion flag
func (c *TCPConnection) GetMarkForDeletion() bool {
	c.RLock()
	defer c.RUnlock()
	return c.MarkForDeletion
}

// SetState is used to setup the state for the TCP connection
func (c *TCPConnection) SetState(state TCPFlowState) {
	c.state = state
}

// Cleanup will provide information when a connection is removed by a timer.
func (c *TCPConnection) Cleanup() {

	c.Lock()
	if !c.expiredConnection && c.state != TCPData {
		c.expiredConnection = true
		c.Context.Counters().IncrementCounter(counters.ErrTCPConnectionsExpired)
	}
	c.Unlock()
}

// SetLoopbackConnection sets LoopbackConnection field.
func (c *TCPConnection) SetLoopbackConnection(isLoopback bool) {
	// Logging information
	c.loopbackConnection = isLoopback
}

// IsLoopbackConnection sets LoopbackConnection field.
func (c *TCPConnection) IsLoopbackConnection() bool {
	// Logging information
	return c.loopbackConnection
}

// NewTCPConnection returns a TCPConnection information struct
func NewTCPConnection(context *pucontext.PUContext, p *packet.Packet) *TCPConnection {

	var initialSeqNumber uint32
	nonce, err := crypto.GenerateRandomBytes(16)
	if err != nil {
		return nil
	}

	// Default tuple in case the packet is nil.
	tuple := &TCPTuple{}
	if p != nil {
		tuple.SourceAddress = p.SourceAddress()
		tuple.DestinationAddress = p.DestinationAddress()
		tuple.SourcePort = p.SourcePort()
		tuple.DestinationPort = p.DestPort()
		initialSeqNumber = p.TCPSequenceNumber()
	}

	return &TCPConnection{
		PingConfig: &PingConfig{},
		state:      TCPSynSend,
		Context:    context,
		Auth: AuthInfo{
			LocalContext: nonce,
		},
		initialSequenceNumber: initialSeqNumber,
		TCPtuple:              tuple,
	}
}

// ProxyConnection is a record to keep state of proxy auth
type ProxyConnection struct {
	sync.Mutex

	state            ProxyConnState
	Auth             AuthInfo
	ReportFlowPolicy *policy.FlowPolicy
	PacketFlowPolicy *policy.FlowPolicy
	reported         bool
	Secrets          secrets.Secrets
}

// NewProxyConnection returns a new Proxy Connection
func NewProxyConnection() *ProxyConnection {
	nonce, err := crypto.GenerateRandomBytes(16)
	if err != nil {
		return nil
	}

	return &ProxyConnection{
		state: ClientTokenSend,
		Auth: AuthInfo{
			LocalContext: nonce,
		},
	}
}

// GetState returns the state of a proxy connection
func (c *ProxyConnection) GetState() ProxyConnState {

	return c.state
}

// SetState is used to setup the state for the Proxy Connection
func (c *ProxyConnection) SetState(state ProxyConnState) {

	c.state = state
}

// SetReported sets the flag to reported when the conn is reported
func (c *ProxyConnection) SetReported(reported bool) {
	c.reported = reported
}

// UDPConnection is information regarding UDP connection.
type UDPConnection struct {
	sync.RWMutex

	state   UDPFlowState
	Context *pucontext.PUContext
	Auth    AuthInfo

	ReportFlowPolicy *policy.FlowPolicy
	PacketFlowPolicy *policy.FlowPolicy
	// ServiceData allows services to associate state with a connection
	ServiceData interface{}

	// PacketQueue indicates app UDP packets queued while authorization is in progress.
	PacketQueue chan *packet.Packet
	Writer      afinetrawsocket.SocketWriter
	// ServiceConnection indicates that this connection is handled by a service
	ServiceConnection bool

	// Stop channels for restransmissions
	synStop    chan bool
	synAckStop chan bool
	ackStop    chan bool

	TestIgnore           bool
	udpQueueFullDropCntr uint64
	expiredConnection    bool

	Secrets secrets.Secrets

	SourceController      string
	DestinationController string
}

// NewUDPConnection returns UDPConnection struct.
func NewUDPConnection(context *pucontext.PUContext, writer afinetrawsocket.SocketWriter) *UDPConnection {

	nonce, err := crypto.GenerateRandomBytes(16)
	if err != nil {
		return nil
	}

	return &UDPConnection{
		state:       UDPStart,
		Context:     context,
		PacketQueue: make(chan *packet.Packet, MaximumUDPQueueLen),
		Writer:      writer,
		Auth: AuthInfo{
			LocalContext: nonce,
		},
		synStop:    make(chan bool),
		synAckStop: make(chan bool),
		ackStop:    make(chan bool),
		TestIgnore: true,
	}
}

// SynStop issues a stop on the synStop channel.
func (c *UDPConnection) SynStop() {
	select {
	case c.synStop <- true:
	default:
		zap.L().Debug("Packet loss - channel was already done")
	}

}

// SynAckStop issues a stop in the synAckStop channel.
func (c *UDPConnection) SynAckStop() {
	select {
	case c.synAckStop <- true:
	default:
		zap.L().Debug("Packet loss - channel was already done")
	}
}

// AckStop issues a stop in the Ack channel.
func (c *UDPConnection) AckStop() {
	select {
	case c.ackStop <- true:
	default:
		zap.L().Debug("Packet loss - channel was already done")
	}

}

// SynChannel returns the SynStop channel.
func (c *UDPConnection) SynChannel() chan bool {
	return c.synStop

}

// SynAckChannel returns the SynAck stop channel.
func (c *UDPConnection) SynAckChannel() chan bool {
	return c.synAckStop
}

// AckChannel returns the Ack stop channel.
func (c *UDPConnection) AckChannel() chan bool {
	return c.ackStop
}

// GetState is used to get state of UDP Connection.
func (c *UDPConnection) GetState() UDPFlowState {
	return c.state
}

// SetState is used to setup the state for the UDP Connection.
func (c *UDPConnection) SetState(state UDPFlowState) {
	c.state = state
}

// QueuePackets queues UDP packets till the flow is authenticated.
func (c *UDPConnection) QueuePackets(udpPacket *packet.Packet) (err error) {
	buffer := make([]byte, len(udpPacket.GetBuffer(0)))
	copy(buffer, udpPacket.GetBuffer(0))

	copyPacket, err := packet.New(packet.PacketTypeApplication, buffer, udpPacket.Mark, true)
	if err != nil {
		return fmt.Errorf("Unable to copy packets to queue:%s", err)
	}

	select {
	case c.PacketQueue <- copyPacket:
	default:
		// connection object is always locked.
		c.udpQueueFullDropCntr++
		return fmt.Errorf("Queue is full")
	}

	return nil
}

// DropPackets drops packets on errors during Authorization.
func (c *UDPConnection) DropPackets() {
	close(c.PacketQueue)
	c.PacketQueue = make(chan *packet.Packet, MaximumUDPQueueLen)
}

// ReadPacket reads a packet from the queue.
func (c *UDPConnection) ReadPacket() *packet.Packet {
	select {
	case p := <-c.PacketQueue:
		return p
	default:
		return nil
	}
}

// Cleanup is called on cache expiry of the connection to record incomplete connections
func (c *UDPConnection) Cleanup() {

	c.Lock()
	if !c.expiredConnection && c.state != UDPData {
		c.expiredConnection = true
		c.Context.Counters().IncrementCounter(counters.ErrUDPConnectionsExpired)
	}
	c.Unlock()
}

// String returns a printable version of connection
func (c *UDPConnection) String() string {

	return fmt.Sprintf("udp-conn state:%d auth: %+v", c.state, c.Auth)
}

// UDPConnectionExpirationNotifier expiration notifier when cache entry expires
func UDPConnectionExpirationNotifier(c cache.DataStore, id interface{}, item interface{}) {

	if conn, ok := item.(*UDPConnection); ok {
		conn.Cleanup()
	}
}