conn.go

// Copyright 2015 Tony Bai.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package cmpp

import (
	"encoding/binary"
	"errors"
	"io"
	"net"
	"sync"
	"time"
)

type State uint8

// Errors for conn operations
var (
	ErrConnIsClosed       = errors.New("connection is closed")
	ErrReadCmdIDTimeout   = errors.New("read commandId timeout")
	ErrReadPktBodyTimeout = errors.New("read packet body timeout")
)

var noDeadline = time.Time{}

// Conn States
const (
	CONN_CLOSED State = iota
	CONN_CONNECTED
	CONN_AUTHOK
)

type Conn struct {
	net.Conn
	State State
	Typ   Type

	// for SeqId generator goroutine
	SeqId <-chan uint32
	done  chan<- struct{}
}

func newSeqIdGenerator() (<-chan uint32, chan<- struct{}) {
	out := make(chan uint32)
	done := make(chan struct{})

	go func() {
		var i uint32
		for {
			select {
			case out <- i:
				i++
			case <-done:
				close(out)
				return
			}
		}
	}()
	return out, done
}

// New returns an abstract structure for successfully
// established underlying net.Conn.
func NewConn(conn net.Conn, typ Type) *Conn {
	seqId, done := newSeqIdGenerator()
	c := &Conn{
		Conn:  conn,
		Typ:   typ,
		SeqId: seqId,
		done:  done,
	}
	tc := c.Conn.(*net.TCPConn) // Always tcpconn
	tc.SetKeepAlive(true)       //Keepalive as default
	return c
}

func (c *Conn) Close() {
	if c != nil {
		if c.State == CONN_CLOSED {
			return
		}
		close(c.done)  // let the SeqId goroutine exit.
		c.Conn.Close() // close the underlying net.Conn
		c.State = CONN_CLOSED
	}
}

func (c *Conn) SetState(state State) {
	c.State = state
}

// SendPkt pack the cmpp packet structure and send it to the other peer.
func (c *Conn) SendPkt(packet Packer, seqId uint32) error {
	if c.State == CONN_CLOSED {
		return ErrConnIsClosed
	}

	data, err := packet.Pack(seqId)
	if err != nil {
		return err
	}

	_, err = c.Conn.Write(data) //block write
	if err != nil {
		return err
	}

	return nil
}

const (
	defaultReadBufferSize = 4096
)

// readBuffer is used to optimize the performance of
// RecvAndUnpackPkt.
type readBuffer struct {
	totalLen  uint32
	commandId CommandId
	leftData  [defaultReadBufferSize]byte
}

var readBufferPool = sync.Pool{
	New: func() interface{} {
		return &readBuffer{}
	},
}

// RecvAndUnpackPkt receives cmpp byte stream, and unpack it to some cmpp packet structure.
func (c *Conn) RecvAndUnpackPkt(timeout time.Duration) (interface{}, error) {
	if c.State == CONN_CLOSED {
		return nil, ErrConnIsClosed
	}
	defer c.SetReadDeadline(noDeadline)

	rb := readBufferPool.Get().(*readBuffer)
	defer readBufferPool.Put(rb)

	// Total_Length in packet
	if timeout != 0 {
		c.SetReadDeadline(time.Now().Add(timeout))
	}
	err := binary.Read(c.Conn, binary.BigEndian, &rb.totalLen)
	if err != nil {
		return nil, err
	}

	if c.Typ == V30 {
		if rb.totalLen < CMPP3_PACKET_MIN || rb.totalLen > CMPP3_PACKET_MAX {
			return nil, ErrTotalLengthInvalid
		}
	}

	if c.Typ == V21 || c.Typ == V20 {
		if rb.totalLen < CMPP2_PACKET_MIN || rb.totalLen > CMPP2_PACKET_MAX {
			return nil, ErrTotalLengthInvalid
		}
	}

	// Command_Id
	if timeout != 0 {
		c.SetReadDeadline(time.Now().Add(timeout))
	}
	err = binary.Read(c.Conn, binary.BigEndian, &rb.commandId)
	if err != nil {
		netErr, ok := err.(net.Error)
		if ok {
			if netErr.Timeout() {
				return nil, ErrReadCmdIDTimeout
			}
		}
		return nil, err
	}

	if !((rb.commandId > CMPP_REQUEST_MIN && rb.commandId < CMPP_REQUEST_MAX) ||
		(rb.commandId > CMPP_RESPONSE_MIN && rb.commandId < CMPP_RESPONSE_MAX)) {
		return nil, ErrCommandIdInvalid
	}

	// The left packet data (start from seqId in header).
	if timeout != 0 {
		c.SetReadDeadline(time.Now().Add(timeout))
	}
	var leftData = rb.leftData[0:(rb.totalLen - 8)]
	_, err = io.ReadFull(c.Conn, leftData)
	if err != nil {
		netErr, ok := err.(net.Error)
		if ok {
			if netErr.Timeout() {
				return nil, ErrReadPktBodyTimeout
			}
		}
		return nil, err
	}

	var p Packer
	switch rb.commandId {
	case CMPP_CONNECT:
		p = &CmppConnReqPkt{}
	case CMPP_CONNECT_RESP:
		if c.Typ == V30 {
			p = &Cmpp3ConnRspPkt{}
		} else {
			p = &Cmpp2ConnRspPkt{}
		}
	case CMPP_TERMINATE:
		p = &CmppTerminateReqPkt{}
	case CMPP_TERMINATE_RESP:
		p = &CmppTerminateRspPkt{}
	case CMPP_SUBMIT:
		if c.Typ == V30 {
			p = &Cmpp3SubmitReqPkt{}
		} else {
			p = &Cmpp2SubmitReqPkt{}
		}
	case CMPP_SUBMIT_RESP:
		if c.Typ == V30 {
			p = &Cmpp3SubmitRspPkt{}
		} else {
			p = &Cmpp2SubmitRspPkt{}
		}
	case CMPP_DELIVER:
		if c.Typ == V30 {
			p = &Cmpp3DeliverReqPkt{}
		} else {
			p = &Cmpp2DeliverReqPkt{}
		}
	case CMPP_DELIVER_RESP:
		if c.Typ == V30 {
			p = &Cmpp3DeliverRspPkt{}
		} else {
			p = &Cmpp2DeliverRspPkt{}
		}
	case CMPP_FWD:
		if c.Typ == V30 {
			p = &Cmpp3FwdReqPkt{}
		} else {
			p = &Cmpp2FwdReqPkt{}
		}
	case CMPP_FWD_RESP:
		if c.Typ == V30 {
			p = &Cmpp3FwdRspPkt{}
		} else {
			p = &Cmpp2FwdRspPkt{}
		}
	case CMPP_ACTIVE_TEST:
		p = &CmppActiveTestReqPkt{}
	case CMPP_ACTIVE_TEST_RESP:
		p = &CmppActiveTestRspPkt{}

	default:
		p = nil
		return nil, ErrCommandIdNotSupported
	}

	err = p.Unpack(leftData)
	if err != nil {
		return nil, err
	}
	return p, nil
}