117 lines
2.2 KiB
Go
117 lines
2.2 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"net"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type PacketQueue struct {
|
|
mu sync.RWMutex
|
|
expectedPackages map[uint32]*Listener
|
|
}
|
|
|
|
type CallResult struct {
|
|
StatusCode uint32
|
|
Data []byte
|
|
}
|
|
|
|
type Listener struct {
|
|
Count int
|
|
Chan chan CallResult
|
|
}
|
|
|
|
func NewPacketQueue(connection net.Conn) *PacketQueue {
|
|
queue := &PacketQueue{
|
|
expectedPackages: make(map[uint32]*Listener),
|
|
}
|
|
go queue.HandleConnection(connection)
|
|
return queue
|
|
}
|
|
|
|
func (p *PacketQueue) RemoveListeners() {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
for _, l := range p.expectedPackages {
|
|
close(l.Chan)
|
|
}
|
|
p.expectedPackages = make(map[uint32]*Listener)
|
|
}
|
|
|
|
func (p *PacketQueue) HandleConnection(connection net.Conn) error {
|
|
defer connection.Close()
|
|
defer p.RemoveListeners()
|
|
var packet Packet
|
|
reader := bufio.NewReader(connection)
|
|
connection.SetReadDeadline(time.Now().Add(time.Millisecond * 200))
|
|
for {
|
|
err := ReadPacket(reader, &packet)
|
|
if err != nil {
|
|
if err == io.EOF {
|
|
return nil
|
|
}
|
|
log.Printf("Error receiving packet: %v\n", err)
|
|
return err
|
|
}
|
|
if packet.Version != CurrentPacketVersion {
|
|
log.Printf("Incorrect packet version: %v\n", packet.Version)
|
|
return fmt.Errorf("incorrect packet version: %d", packet.Version)
|
|
}
|
|
if packet.DataLength == 0 {
|
|
go p.HandleData(packet.MessageType, CallResult{
|
|
StatusCode: packet.StatusCode,
|
|
Data: []byte{},
|
|
})
|
|
continue
|
|
}
|
|
data, err := GetPacketData(reader, packet.DataLength)
|
|
if err != nil {
|
|
log.Printf("Error receiving packet data: %v\n", err)
|
|
return err
|
|
} else {
|
|
go p.HandleData(packet.MessageType, CallResult{
|
|
StatusCode: packet.StatusCode,
|
|
Data: data,
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func (p *PacketQueue) HandleData(t uint32, data CallResult) {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
l, ok := p.expectedPackages[t]
|
|
if ok {
|
|
l.Chan <- data
|
|
l.Count--
|
|
if l.Count == 0 {
|
|
close(l.Chan)
|
|
delete(p.expectedPackages, t)
|
|
}
|
|
return
|
|
}
|
|
}
|
|
|
|
func (p *PacketQueue) Expect(messageType uint32) <-chan CallResult {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
l, ok := p.expectedPackages[messageType]
|
|
if ok {
|
|
l.Count++
|
|
return l.Chan
|
|
}
|
|
|
|
ch := make(chan CallResult)
|
|
p.expectedPackages[messageType] = &Listener{
|
|
Count: 1,
|
|
Chan: ch,
|
|
}
|
|
|
|
return ch
|
|
|
|
}
|