/
router.go
95 lines (78 loc) · 2.13 KB
/
router.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
// Copyright 2023 FishGoddess. All rights reserved.
// Use of this source code is governed by a MIT style
// license that can be found in the LICENSE file.
package pack
import (
"context"
"errors"
"io"
"sync"
"github.com/FishGoddess/vex"
"github.com/FishGoddess/vex/log"
)
var (
errPacketHandlerNotFound = errors.New("vex: packet handler not found")
)
// PacketHandler is a handler for handling packets.
// You will receive a byte slice of request and should return a byte slice or error if necessary.
type PacketHandler func(ctx context.Context, packetType PacketType, requestPacket []byte) (responsePacket []byte, err error)
type Router struct {
handlers map[PacketType]PacketHandler
lock sync.RWMutex
}
// NewRouter creates a router for registering some packet handlers.
func NewRouter() *Router {
return &Router{
handlers: make(map[PacketType]PacketHandler, 16),
}
}
// Register registers a packet handler to router.
func (r *Router) Register(packetType PacketType, handler PacketHandler) {
r.lock.Lock()
r.handlers[packetType] = handler
r.lock.Unlock()
}
func (r *Router) writeStandardPacket(writer io.Writer, data []byte) {
err := writePacket(writer, packetTypeStandard, data)
if err != nil {
log.Error(err, "write standard packet failed")
}
}
func (r *Router) writeErrorPacket(writer io.Writer, err error) {
err = writePacket(writer, packetTypeError, []byte(err.Error()))
if err != nil {
log.Error(err, "write error packet failed")
}
}
// Handle handles context of vex, and you can pass it to a server.
func (r *Router) Handle(ctx *vex.Context) {
for {
select {
case <-ctx.Done():
log.Debug("router context of %s has done", ctx.RemoteAddr())
return
default:
}
packetType, data, err := readPacket(ctx)
if err == io.EOF {
return
}
if err != nil {
r.writeErrorPacket(ctx, err)
continue
}
r.lock.RLock()
handle, ok := r.handlers[packetType]
r.lock.RUnlock()
if !ok {
r.writeErrorPacket(ctx, errPacketHandlerNotFound)
continue
}
data, err = handle(ctx, packetType, data)
if err != nil {
r.writeErrorPacket(ctx, err)
continue
}
r.writeStandardPacket(ctx, data)
}
}