package beeq import ( "encoding/binary" "github.com/zgwit/beeq/packet" "log" "net" "time" ) type Bee struct { clientId string session *Session conn net.Conn hive *Hive quit chan struct{} events chan *Event timeout time.Duration alive bool } func NewBee(conn net.Conn, hive *Hive) *Bee { return &Bee{ conn: conn, hive: hive, quit: make(chan struct{}), events: make(chan *Event, 1024), timeout: time.Hour * 24, } } func (bee *Bee) Event(event *Event) { bee.events <- event } func (bee *Bee) Active() { bee.alive = true go bee.receiver() go bee.messenger() } func (bee *Bee) Shutdown() { bee.alive = false // Tell hive delete me bee.hive.Event(NewEvent(E_LOST_CONN, nil, bee)) // Stop messenger close(bee.events) close(bee.quit) } func (bee *Bee) recv(b []byte) (int, error) { err := bee.conn.SetReadDeadline(time.Now().Add(bee.timeout)) if err != nil { return 0, err } return bee.conn.Read(b) } func (bee *Bee) send(b []byte) (int, error) { err := bee.conn.SetWriteDeadline(time.Now().Add(bee.timeout)) if err != nil { return 0, err } return bee.conn.Write(b) } func (bee *Bee) receiver() { //Abort error defer func() { if r := recover(); r!=nil { log.Print("bee receiver panic ", r) bee.Shutdown() } }() readHead := true buf := Alloc(6) offset := 0 total := 0 for { if l, err := bee.recv(buf[offset:]); err != nil { log.Print("Receive Failed: ", err) break } else { offset += l } //Parse head if readHead && offset >= 2 { for i := 1; i <= offset; i++ { if buf[i] < 0x80 { rl, rll := binary.Uvarint(buf[1:]) remainLength := int(rl) total = remainLength + rll + 1 if total > 6 { buf = ReAlloc(buf, total) } readHead = false break } } } //Parse Message if !readHead && offset >= total { readHead = true b := Alloc(6) if msg, l, err := packet.Decode(buf); err != nil { //TODO log err log.Print(err) offset = 0 //clear data } else { bee.Event(NewEvent(E_MESSAGE, msg, bee)) //Only message less than 6 bytes if offset > l { copy(b, buf[l:]) offset -= l } else { offset = 0 //clear data } } buf = b } } //Shutdown bee if bee.alive { bee.Shutdown() } } func (bee *Bee) messenger() { //Abort error defer func() { if r := recover(); r!=nil { log.Print("bee messenger panic ", r) bee.Shutdown() } }() for { //Blocking var event *Event select { case <-bee.quit: bee.conn.Close() return case event = <-bee.events: switch event.event { case E_MESSAGE: bee.handleMessage(event.data.(packet.Message)) case E_DISPATCH: bee.dispatchMessage(event.data.(packet.Message)) case E_CLOSE: bee.Shutdown() } } } } func (bee *Bee) handleMessage(msg packet.Message) { log.Printf("Received message from %s: %s QOS(%d) DUP(%t) RETAIN(%t)", bee.clientId, msg.Type().Name(), msg.Qos(), msg.Dup(), msg.Retain()) //log.Print("recv Message:", msg.Type().Name()) switch msg.Type() { case packet.CONNECT: bee.handleConnect(msg.(*packet.Connect)) case packet.PUBLISH: bee.handlePublish(msg.(*packet.Publish)) case packet.PUBACK: bee.handlePubAck(msg.(*packet.PubAck)) case packet.PUBREC: bee.handlePubRec(msg.(*packet.PubRec)) case packet.PUBREL: bee.handlePubRel(msg.(*packet.PubRel)) case packet.PUBCOMP: bee.handlePubComp(msg.(*packet.PubComp)) case packet.SUBSCRIBE: bee.handleSubscribe(msg.(*packet.Subscribe)) case packet.UNSUBSCRIBE: bee.handleUnSubscribe(msg.(*packet.UnSubscribe)) case packet.PINGREQ: bee.handlePingReq(msg.(*packet.PingReq)) case packet.DISCONNECT: bee.handleDisconnect(msg.(*packet.DisConnect)) } } func (bee *Bee) handleConnect(msg *packet.Connect) { bee.clientId = string(msg.ClientId()) //if msg.WillFlag() { // bee.session.will = new(packet.Publish) //} bee.hive.Event(NewEvent(E_CONNECT, msg, bee)) } func (bee *Bee) handlePublish(msg *packet.Publish) { qos := msg.Qos() if qos == packet.Qos0 { bee.hive.Event(NewEvent(E_PUBLISH, msg, bee)) } else if qos == packet.Qos1 { bee.hive.Event(NewEvent(E_PUBLISH, msg, bee)) //Reply PUBACK puback := packet.PUBACK.NewMessage().(*packet.PubAck) puback.SetPacketId(msg.PacketId()) bee.Event(NewEvent(E_DISPATCH, puback, bee)) } else if qos == packet.Qos2 { //Save & Send PUBREC bee.session.recvPub2[msg.PacketId()] = msg pubrec := packet.PUBREC.NewMessage().(*packet.PubRec) pubrec.SetPacketId(msg.PacketId()) bee.Event(NewEvent(E_DISPATCH, pubrec, bee)) } else { //error } } func (bee *Bee) handlePubAck(msg *packet.PubAck) { if _, ok := bee.session.pub1[msg.PacketId()]; ok { delete(bee.session.pub1, msg.PacketId()) } } func (bee *Bee) handlePubRec(msg *packet.PubRec) { msg.SetType(packet.PUBREL) bee.Event(NewEvent(E_DISPATCH, msg, bee)) } func (bee *Bee) handlePubRel(msg *packet.PubRel) { msg.SetType(packet.PUBCOMP) bee.Event(NewEvent(E_DISPATCH, msg, bee)) } func (bee *Bee) handlePubComp(msg *packet.PubComp) { if _, ok := bee.session.pub2[msg.PacketId()]; ok { delete(bee.session.pub2, msg.PacketId()) } } func (bee *Bee) handleSubscribe(msg *packet.Subscribe) { bee.hive.Event(NewEvent(E_SUBSCRIBE, msg, bee)) } func (bee *Bee) handleUnSubscribe(msg *packet.UnSubscribe) { bee.hive.Event(NewEvent(E_UNSUBSCRIBE, msg, bee)) } func (bee *Bee) handlePingReq(msg *packet.PingReq) { msg.SetType(packet.PINGRESP) bee.Event(NewEvent(E_DISPATCH, msg, bee)) } func (bee *Bee) handleDisconnect(msg *packet.DisConnect) { bee.hive.Event(NewEvent(E_DISCONNECT, msg, bee)) } func (bee *Bee) dispatchMessage(msg packet.Message) { log.Printf("Send message to %s: %s QOS(%d) DUP(%t) RETAIN(%t)", bee.clientId, msg.Type().Name(), msg.Qos(), msg.Dup(), msg.Retain()) if head, payload, err := msg.Encode(); err != nil { //TODO log log.Print("Message encode error: ", err) } else { bee.send(head) if payload != nil && len(payload) > 0 { bee.send(payload) } } if msg.Type() == packet.PUBLISH { pub := msg.(*packet.Publish) //Publish Qos1 Qos2 Need store if msg.Qos() == packet.Qos1 { bee.session.pub1[pub.PacketId()] = pub } else if msg.Qos() == packet.Qos2 { bee.session.pub2[pub.PacketId()] = pub } } }